Skip to content

Instantly share code, notes, and snippets.

@jozefg
Last active August 3, 2017 13:55
Show Gist options
  • Save jozefg/4c7aa55dc6566a5ac4c6 to your computer and use it in GitHub Desktop.
Save jozefg/4c7aa55dc6566a5ac4c6 to your computer and use it in GitHub Desktop.
A simple trick to pick a random element from a stream in constant memory
{-# LANGUAGE FlexibleContexts #-}
module PickRandom where
import Data.List (group, sort)
import Control.Monad
import Control.Monad.Random (MonadRandom, getRandomR)
import qualified Control.Foldl as F
-- Pick a value uniformly from a fold
pickRandom :: MonadRandom m => a -> F.FoldM m a a
pickRandom a = F.FoldM choose (return (a, 0 :: Int)) (return . fst)
where choose (current, idx) new = do
let idx' = idx + 1
shouldTake <- (== 0) <$> getRandomR (0, idx')
if shouldTake
then return (new, idx')
else return (current, idx')
-- As a quick demonstration, pick 10k elements and print distribution
main :: IO ()
main = process <$> replicateM 10000 pick >>= print
where pick = F.foldM (pickRandom (-1)) [0..10]
process = map length . group . sort
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment