リストのシャッフル

http://d.hatena.ne.jp/yts/20051022#p1 で言及されたアルゴリズムを実装してみる。

{-# OPTIONS -fglasgow-exts #-}
module Shuffle where
import Control.Monad.Fix (fix)
import System.Random

import Control.Monad.ST

import Data.Array.MArray
import Data.Array.ST

シャッフルのための互換の列を生成する関数。

swaps :: (Int -> Int -> a) -> StdGen -> [a]
swaps f gen = map fst . fix $ \lst -> 
    (f 0 0, (0, gen)):[(f j a, (j, g')) | (_, (i, g)) <- lst
	, let j = i + 1; (a, g') = randomR (0, j) g]
*Shuffle> take 20 $ swaps (,) (mkStdGen 0)
[(0,0),(1,1),(2,2),(3,1),(4,0),(5,4),(6,3),(7,3),(8,1),(9,3),(10,4),(11,11),(12,4),(13,8),(14,7),(15,8),(16,14),(17,2),(18,14),(19,7)]

互換の効率のためにリストをSTArrayに変換してからシャッフルする。

listShuffle :: StdGen -> [e] -> [e]
listShuffle gen (lst :: [e]) = runST (
	  foldl (>>=)
	    (newListArray (0, len - 1) lst :: ST s (STArray s Int e))
	    (take len $ swaps mArraySwap gen)
	  >>= getElems)
  where
    len = length lst

mArraySwap i j ar
  = do	ei <- readArray ar i
	ej <- readArray ar j
	writeArray ar j ei
	writeArray ar i ej
	return ar