Personal tools

Random shuffle

From HaskellWiki

(Difference between revisions)
Jump to: navigation, search
(Random shuffle algorithm)
 
(Purely functional)
 
(6 intermediate revisions by 5 users not shown)
Line 8: Line 8:
   
 
<haskell>
 
<haskell>
{-# LANGUAGE ScopedTypeVariables #-}
 
 
 
import System.Random
 
import System.Random
 
import Data.Array.IO
 
import Data.Array.IO
Line 16: Line 14:
 
-- | Randomly shuffle a list
 
-- | Randomly shuffle a list
 
-- /O(N)/
 
-- /O(N)/
shuffle :: forall a. [a] -> IO [a]
+
shuffle :: [a] -> IO [a]
 
shuffle xs = do
 
shuffle xs = do
let n = length xs
+
ar <- newArray n xs
ar <- newListArray (1,n) xs :: IO (IOArray Int a)
 
 
forM [1..n] $ \i -> do
 
forM [1..n] $ \i -> do
 
j <- randomRIO (i,n)
 
j <- randomRIO (i,n)
Line 25: Line 23:
 
writeArray ar j vi
 
writeArray ar j vi
 
return vj
 
return vj
  +
where
  +
n = length xs
  +
newArray :: Int -> [a] -> IO (IOArray Int a)
  +
newArray n xs = newListArray (1,n) xs
  +
</haskell>
  +
  +
Or one can use ST to avoid needing IO:
  +
  +
<haskell>
  +
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
  +
</haskell>
  +
  +
And if you are using IO's hidden StdGen you can wrap this as usual:
  +
  +
<haskell>
  +
shuffleIO :: [a] -> IO [a]
  +
shuffleIO xs = getStdRandom (shuffle' xs)
 
</haskell>
 
</haskell>
   
 
This is a lot simpler than the purely functional algorithm linked below.
 
This is a lot simpler than the purely functional algorithm linked below.
  +
  +
Here's a variation using the MonadRandom package:
  +
  +
<haskell>
  +
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]
  +
</haskell>
   
 
== Other implemenations ==
 
== Other implemenations ==
   
  +
=== Purely functional ===
  +
* Using Data.Map, O(n * log n)
  +
<haskell>
  +
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)
  +
</haskell>
 
* [http://okmij.org/ftp/Haskell/perfect-shuffle.txt Purely functional O(n log n) random shuffle algorithm].
 
* [http://okmij.org/ftp/Haskell/perfect-shuffle.txt Purely functional O(n log n) random shuffle algorithm].
  +
  +
=== 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
  +
  +
<haskell>
  +
{- | @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
  +
</haskell>
  +
  +
So e.g.
  +
<haskell>
  +
*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"]
  +
</haskell>

Latest revision as of 19:28, 11 January 2011

Contents

[edit] 1 The problem

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

[edit] 2 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]

[edit] 3 Other implemenations

[edit] 3.1 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)

[edit] 3.2 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"]