Random shuffle

From HaskellWiki
Jump to navigation Jump to search
The printable version is no longer supported and may have rendering errors. Please update your browser bookmarks and please use the default browser print function instead.

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:

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

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

Or one can use ST to avoid needing IO:

import System.Random
import Data.Array.ST
import Control.Monad
import Control.Monad.ST
import Data.STRef

-- | 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'))
  where
    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.

Here's a variation using the MonadRandom package:

import Control.Monad
import Control.Monad.ST
import Control.Monad.Random
import System.Random
import Data.Array.ST
import GHC.Arr

shuffle :: RandomGen g => [a] -> Rand g [a]
shuffle xs = do
    let l = length xs
    rands <- take l `fmap` getRandomRs (0, l-1)
    let ar = runSTArray $ do
        ar <- thawSTArray $ listArray (0, l-1) xs
        forM_ (zip [0..(l-1)] rands) $ \(i, j) -> do
            vi <- readSTArray ar i
            vj <- readSTArray ar j
            writeSTArray ar j vi
            writeSTArray ar i vj
        return ar
    return (elems ar)

*Main> evalRandIO (shuffle [1..10])
[6,5,1,7,10,4,9,2,8,3]

Other implemenations

Purely functional

  • Using Data.Map, O(n * log n)
import System.Random
import Data.Map

fisherYatesStep :: RandomGen g => (Map Int a, g) -> (Int, a) -> (Map Int a, g)
fisherYatesStep (m, gen) (i, x) = ((insert j x . insert i (m ! j)) m, gen')
  where
    (j, gen') = randomR (0, i) gen

fisherYates :: RandomGen g => g -> [a] -> ([a], g)
fisherYates gen [] = ([], gen)
fisherYates gen l = 
  toElems $ foldl fisherYatesStep (initial (head l) gen) (numerate (tail l))
  where
    toElems (x, y) = (elems x, y)
    numerate = zip [1..]
    initial x gen = (singleton 0 x, gen)

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
    where
        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)
    where
        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)
["fbd","efb","bef","adc","cef","eac"]
*Main MonadRandom Random> grabble "abcdef" 6 3
["fce","dfa","ebf","edb","cea","dbc"]
*Main MonadRandom Random> grabble "abcdef" 6 3
["cbf","dec","edb","fae","bda","cde"]