RosettaCodeData/Task/Knuths-algorithm-S/Haskell/knuths-algorithm-s.hs

49 lines
1.2 KiB
Haskell

import Control.Monad.Random
import Control.Monad.State
import qualified Data.Map as M
import System.Random
-- s_of_n_creator :: Int -> a -> RandT StdGen (State (Int, [a])) [a]
s_of_n_creator :: Int -> a -> StateT (Int, [a]) (Rand StdGen) [a]
s_of_n_creator n v = do
(i, vs) <- get
let i' = i + 1
if i' <= n
then do
let vs' = v : vs
put (i', vs')
pure vs'
else do
j <- getRandomR (1, i')
if j > n
then do
put (i', vs)
pure vs
else do
k <- getRandomR (0, n - 1)
let (f, (_ : b)) = splitAt k vs
vs' = v : f ++ b
put (i', vs')
pure vs'
sample :: Int -> Rand StdGen [Int]
sample n =
let s_of_n = s_of_n_creator n
in snd <$> execStateT (traverse s_of_n [0 .. 9 :: Int]) (0, [])
incEach :: (Ord a, Num b) => M.Map a b -> [a] -> M.Map a b
incEach m ks = foldl (\m' k -> M.insertWith (+) k 1 m') m ks
sampleInc :: Int -> M.Map Int Double -> Rand StdGen (M.Map Int Double)
sampleInc n m = do
s <- sample n
pure $ incEach m s
main :: IO ()
main = do
let counts = M.empty :: M.Map Int Double
n = 100000
gen <- getStdGen
counts <- evalRandIO $ foldM (\c _ -> sampleInc 3 c) M.empty [1 .. n]
print (fmap (/ n) counts)