Search code examples
algorithmhaskellfunctional-programmingshufflepurely-functional

pure Knuth/Fisher-Yates shuffle in haskell


In go I could write a function like this:

func pureFisherYates(s []int, swaps []int) []int {
    newS := copy(s)
    for i, _ := range newS {
            for _, j := range swaps {
                    newS[i], newS[j] = newS[j], newS[i]
            }
    }
}

To me this seems like a pure function. It always returns the same output given the same input, and it doesn't mutate the state of the world (except in some strict sense the same way that any other function does, taking up cpu resources, creating heat energy, etc). Yet whenever I look for how to do pure shuffling I find stuff like this, and whenever I look for specifically Haskell implementation Fisher-Yates I either get an 0^2 Fisher-Yates implemented with a list or a [a] -> IO [a] implementation. Does there exist a [a] -> [a] O(n) shuffle and if not why is my above go implementation impure.


Solution

  • The ST monad allows exactly such encapsulated mutability, and Data.Array.ST contains arrays which can be mutated in ST and then an immutable version returned outside.

    https://wiki.haskell.org/Random_shuffle gives two implementations of Fisher-Yates shuffle using ST. They aren't literally [a] -> [a], but that's because random number generation needs to be handled as well:

    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

    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 <- forM [0..(l-2)] $ \i -> getRandomR (i, l-1)
        let ar = runSTArray $ do
            ar <- thawSTArray $ listArray (0, l-1) xs
            forM_ (zip [0..] 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: with a fixed swaps argument as in your Go code, the code is quite simple

    {-# LANGUAGE ScopedTypeVariables #-}
    
    import Data.Array.ST
    import Data.Foldable
    import Control.Monad.ST
    
    shuffle :: forall a. [a] -> [Int] -> [a]
    shuffle xs swaps = runST $ do
        let n = length xs
        ar <- newListArray (1,n) xs :: ST s (STArray s Int a)
        for_ [1..n] $ \i ->
            for_ swaps $ \j -> do
                vi <- readArray ar i
                vj <- readArray ar j
                writeArray ar j vi
                writeArray ar i vj
        getElems ar
    

    but I am not sure you can reasonably call it Fisher-Yates shuffle.