Random shuffle

From HaskellWiki
Revision as of 15:21, 26 October 2007 by Remi (talk | contribs) (drawing without replacement)

Jump to: navigation, search

The problem

Shuffling a list, i.e. creating a random permutation, is not easy to do correctly. Each permutation should have the same probability.

Imperative algorithm

The standard imperative algorithm can be implemented as follows:

{-# LANGUAGE ScopedTypeVariables #-}

import System.Random
import Data.Array.IO
import Control.Monad

-- | Randomly shuffle a list
--   /O(N)/
shuffle :: forall a. [a] -> IO [a]
shuffle xs = do
        let n = length xs
        ar <- newListArray (1,n) xs :: IO (IOArray Int a)
        forM [1..n] $ \i -> do
            j <- randomRIO (i,n)
            vi <- readArray ar i
            vj <- readArray ar j
            writeArray ar j vi
            return vj

Or one can use ST to avoid needing IO:

-- | Randomly shuffle a list without the IO Monad
--   /O(N)/
shuffle' :: [a] -> StdGen -> ([a],StdGen)
shuffle' xs gen = runST (do
        g <- newSTRef gen
        let randomRST lohi = do
              (a,s') <- liftM (randomR lohi) (readSTRef g)
              writeSTRef g s'
              return a
        ar <- newArray n xs
        xs' <- forM [1..n] $ \i -> do
                j <- randomRST (i,n)
                vi <- readArray ar i
                vj <- readArray ar j
                writeArray ar j vi
                return vj
        gen' <- readSTRef g
        return (xs',gen'))
    n = length xs
    newArray :: Int -> [a] -> ST s (STArray s Int a)
    newArray n xs =  newListArray (1,n) xs

And if you are using IO's hidden StdGen you can wrap this as usual:

shuffleIO :: [a] -> IO [a]
shuffleIO xs = getStdRandom (shuffle' xs)

This is a lot simpler than the purely functional algorithm linked below.

Other implemenations

Purely functional

Drawing without replacement

  • uses New_monads/MonadRandom
  • allows you to not shuffle the entire list but only part of it (drawing elements without replacement)
  • allows you to take multiple drawings/shufflings at once, which can save some array building
{- | @grabble xs m n@ is /O(m*n')/, where @n' = min n (length xs)@
     Chooses @n@ elements from @xs@, without putting back,
     and that @m@ times. -}
grabble :: MonadRandom m => [a] -> Int -> Int -> m [[a]]
grabble xs m n = do
    swapss <- replicateM m $ forM [0 .. min (maxIx - 1) n] $ \i -> do
                j <- getRandomR (i, maxIx)
                return (i, j)
    return $ map (take n . swapElems xs) swapss
        maxIx   = length xs - 1

grabbleOnce :: MonadRandom m => [a] -> Int -> m [a]
grabbleOnce xs n = head `liftM` grabble xs 1 n

swapElems  :: [a] -> [(Int, Int)] -> [a]
swapElems xs swaps = elems $ runSTArray (do
    arr <- newListArray (0, maxIx) xs
    mapM_ (swap arr) swaps
    return arr)
        maxIx   = length xs - 1
        swap arr (i,j) = do
            vi <- readArray arr i
            vj <- readArray arr j
            writeArray arr i vj
            writeArray arr j vi

So e.g.

*Main MonadRandom Random> evalRand (grabble "abcdef" 6 3) (mkStdGen 0)
*Main MonadRandom Random> grabble "abcdef" 6 3
*Main MonadRandom Random> grabble "abcdef" 6 3