summaryrefslogtreecommitdiff
path: root/System/Random/Shuffle.hs
diff options
context:
space:
mode:
authorBen Sima <ben@bsima.me>2020-04-15 10:06:24 -0700
committerBen Sima <ben@bsima.me>2020-04-15 10:07:02 -0700
commitafa9d701538b9e56622a0bfdf8e04aa358c9cd82 (patch)
treedee95c3955b3fe3d11e80d89823660d28eee0587 /System/Random/Shuffle.hs
parentf4b8c0df041b063c0b47d2ec6c818a9c202fd833 (diff)
Reformatting
Now I'm using ormolu instead of brittany for Haskell formatting now. Figured I should just make all of these big changes at once.
Diffstat (limited to 'System/Random/Shuffle.hs')
-rw-r--r--System/Random/Shuffle.hs196
1 files changed, 96 insertions, 100 deletions
diff --git a/System/Random/Shuffle.hs b/System/Random/Shuffle.hs
index 02cd3e0..774e7b4 100644
--- a/System/Random/Shuffle.hs
+++ b/System/Random/Shuffle.hs
@@ -1,122 +1,118 @@
-{- |
-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
+-- 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
module System.Random.Shuffle
- ( shuffle
- , shuffle'
- , shuffleM
+ ( shuffle,
+ shuffle',
+ shuffleM,
)
where
-import Data.Function ( fix )
-import System.Random ( RandomGen
- , randomR
- )
-import Control.Monad ( liftM
- , liftM2
- )
-import Control.Monad.Random ( MonadRandom
- , getRandomR
- )
-
+import Control.Monad
+ ( liftM,
+ liftM2,
+ )
+import Control.Monad.Random
+ ( MonadRandom,
+ getRandomR,
+ )
+import Data.Function (fix)
+import System.Random
+ ( RandomGen,
+ randomR,
+ )
-- | 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
-
+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.
+ 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.
+ 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
+ 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))
+ | 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))