From c843083647edc9bfba41e74b83d8cea54a1ac646 Mon Sep 17 00:00:00 2001
From: Ben Sima <ben@bsima.me>
Date: Sat, 23 Feb 2019 22:08:08 -0800
Subject: Add System.Random.Shuffle
---
lore/System/Random/Shuffle.hs | 122 ++++++++++++++++++++++++++++++++++++++++++
1 file changed, 122 insertions(+)
create mode 100644 lore/System/Random/Shuffle.hs
diff --git a/lore/System/Random/Shuffle.hs b/lore/System/Random/Shuffle.hs
new file mode 100644
index 0000000..02cd3e0
--- /dev/null
+++ b/lore/System/Random/Shuffle.hs
@@ -0,0 +1,122 @@
+{- |
+Module : System.Random.Shuffle
+Copyright : (c) 2009 Oleg Kiselyov, Manlio Perillo
+License : BSD3 (see LICENSE file)
+
+<http://okmij.org/ftp/Haskell/perfect-shuffle.txt>
+
+
+Example:
+
+ import System.Random (newStdGen)
+ import System.Random.Shuffle (shuffle')
+
+ main = do
+ rng <- newStdGen
+ let xs = [1,2,3,4,5]
+ print $ shuffle' xs (length xs) rng
+-}
+{-# OPTIONS_GHC -funbox-strict-fields #-}
+
+module System.Random.Shuffle
+ ( shuffle
+ , shuffle'
+ , shuffleM
+ )
+where
+
+import Data.Function ( fix )
+import System.Random ( RandomGen
+ , randomR
+ )
+import Control.Monad ( liftM
+ , liftM2
+ )
+import Control.Monad.Random ( MonadRandom
+ , getRandomR
+ )
+
+
+-- | A complete binary tree, of leaves and internal nodes.
+-- Internal node: Node card l r
+-- where card is the number of leaves under the node.
+-- Invariant: card >=2. All internal tree nodes are always full.
+data Tree a = Leaf !a
+ | Node !Int !(Tree a) !(Tree a)
+ deriving Show
+
+
+-- | Convert a sequence (e1...en) to a complete binary tree
+buildTree :: [a] -> Tree a
+buildTree = (fix growLevel) . (map Leaf)
+ where
+ growLevel _ [node] = node
+ growLevel self l = self $ inner l
+
+ inner [] = []
+ inner [e ] = [e]
+ inner (e1 : e2 : es) = e1 `seq` e2 `seq` (join e1 e2) : inner es
+
+ join l@(Leaf _ ) r@(Leaf _ ) = Node 2 l r
+ join l@(Node ct _ _ ) r@(Leaf _ ) = Node (ct + 1) l r
+ join l@(Leaf _ ) r@(Node ct _ _) = Node (ct + 1) l r
+ join l@(Node ctl _ _) r@(Node ctr _ _) = Node (ctl + ctr) l r
+
+
+-- |Given a sequence (e1,...en) to shuffle, and a sequence
+-- (r1,...r[n-1]) of numbers such that r[i] is an independent sample
+-- from a uniform random distribution [0..n-i], compute the
+-- corresponding permutation of the input sequence.
+shuffle :: [a] -> [Int] -> [a]
+shuffle elements = shuffleTree (buildTree elements)
+ where
+ shuffleTree (Leaf e) [] = [e]
+ shuffleTree tree (r : rs) =
+ let (b, rest) = extractTree r tree in b : (shuffleTree rest rs)
+ shuffleTree _ _ = error "[shuffle] called with lists of different lengths"
+
+ -- Extracts the n-th element from the tree and returns
+ -- that element, paired with a tree with the element
+ -- deleted.
+ -- The function maintains the invariant of the completeness
+ -- of the tree: all internal nodes are always full.
+ extractTree 0 (Node _ (Leaf e) r ) = (e, r)
+ extractTree 1 (Node 2 (Leaf l) (Leaf r)) = (r, Leaf l)
+ extractTree n (Node c (Leaf l) r) =
+ let (e, r') = extractTree (n - 1) r in (e, Node (c - 1) (Leaf l) r')
+
+ extractTree n (Node n' l (Leaf e)) | n + 1 == n' = (e, l)
+
+ extractTree n (Node c l@(Node cl _ _) r)
+ | n < cl
+ = let (e, l') = extractTree n l in (e, Node (c - 1) l' r)
+ | otherwise
+ = let (e, r') = extractTree (n - cl) r in (e, Node (c - 1) l r')
+ extractTree _ _ = error "[extractTree] impossible"
+
+-- |Given a sequence (e1,...en) to shuffle, its length, and a random
+-- generator, compute the corresponding permutation of the input
+-- sequence.
+shuffle' :: RandomGen gen => [a] -> Int -> gen -> [a]
+shuffle' elements len = shuffle elements . rseq len
+ where
+ -- The sequence (r1,...r[n-1]) of numbers such that r[i] is an
+ -- independent sample from a uniform random distribution
+ -- [0..n-i]
+ rseq :: RandomGen gen => Int -> gen -> [Int]
+ rseq n = fst . unzip . rseq' (n - 1)
+ where
+ rseq' :: RandomGen gen => Int -> gen -> [(Int, gen)]
+ rseq' 0 _ = []
+ rseq' i gen = (j, gen) : rseq' (i - 1) gen'
+ where (j, gen') = randomR (0, i) gen
+
+-- |shuffle' wrapped in a random monad
+shuffleM :: (MonadRandom m) => [a] -> m [a]
+shuffleM elements
+ | null elements = return []
+ | otherwise = liftM (shuffle elements) (rseqM (length elements - 1))
+ where
+ rseqM :: (MonadRandom m) => Int -> m [Int]
+ rseqM 0 = return []
+ rseqM i = liftM2 (:) (getRandomR (0, i)) (rseqM (i - 1))
--
cgit v1.2.3