A simple trick to pick a random element from a stream in constant memory

PickRandom.hs

{-# 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