Difference between revisions of "Random shuffle"

From HaskellWiki
Jump to navigation Jump to search
(Random shuffle algorithm)
 
(Add pure version using ST)
Line 26: Line 26:
 
writeArray ar j vi
 
writeArray ar j vi
 
return vj
 
return vj
  +
</haskell>
  +
  +
Or one can use ST to avoid needing IO:
  +
  +
<haskell>
  +
-- | 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>
   

Revision as of 20:49, 18 October 2007

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

Other implemenations