module Mcmc.Tools.Shuffle
( shuffle,
shuffleN,
grabble,
)
where
import Control.Monad
import Control.Monad.ST
import qualified Data.Vector as V
import Data.Vector (Vector)
import qualified Data.Vector.Mutable as M
import System.Random.MWC
( GenIO,
uniformR,
)
shuffle :: [a] -> GenIO -> IO [a]
shuffle xs g = head <$> grabble xs 1 (length xs) g
shuffleN :: [a] -> Int -> GenIO -> IO [[a]]
shuffleN xs n = grabble xs n (length xs)
grabble :: [a] -> Int -> Int -> GenIO -> IO [[a]]
grabble xs m n gen = do
swapss <- replicateM m $ forM [0 .. min (l - 1) n] $ \i -> do
j <- uniformR (i, l) gen
return (i, j)
return $ map (V.toList . V.take n . swapElems (V.fromList xs)) swapss
where
l = length xs - 1
swapElems :: Vector a -> [(Int, Int)] -> Vector a
swapElems xs swaps = runST $ do
mxs <- V.unsafeThaw xs
mapM_ (uncurry $ M.unsafeSwap mxs) swaps
V.unsafeFreeze mxs