Persistent data structure for random selection without replacement in Haskell.

Data.WithoutReplacement.hs

------------------------------------------------------------------------------
-- |
-- Module       : Data.WithoutReplacement
-- Copyright    : (c) Guilherme G. Brandt, 2022
-- License      : Zlib
-- 
-- Persistent data structure for random selection without replacement.
--
------------------------------------------------------------------------------

module Data.WithoutReplacement(
    WithoutReplacement
    , null
    , size
    , fromList
    , toList
    , toListStd
    , randomWithoutReplacement
    , randomWithoutReplacementM
) where

import System.Random (
    RandomGen
    , getStdGen
    , randomR
    )
import System.Random.Stateful (
    STGenM
    , RandomGenM
    , runSTGen
    , runSTGen_
    , randomRM
    )
import Control.Monad.ST (ST)
import Control.Monad.IO.Class (MonadIO)
import Prelude hiding (null)

-- | Persistent data structure for random selection without replacement.
--
-- This data structure is persistent the sense that it is immutable (as things
-- usually are in Haskell), and the selection is without replacement in that
-- calling 'randomWithoutReplacement' on a version and then calling it again
-- on the new version that is returned will return the same value only as many
-- times as it appears in the originating list.

data WithoutReplacement a =
    Node Int (WithoutReplacement a) (WithoutReplacement a)
    | Leaf a
    | Nil
    deriving (Show, Eq)

{-----------------------------------------------------------------------------
  Bookkeeping functions
------------------------------------------------------------------------------}

-- | /O(1)/. Is the set empty?
--
-- > Data.WithoutReplacement.null (fromList [])         == True
-- > Data.WithoutReplacement.null (fromList [1, 2, 3])  == False

null :: WithoutReplacement a -> Bool
null Nil = True
null _ = False

-- | /O(1)/. The number of elements still in the set.
--
-- > size (fromList [])         == 0
-- > size (fromList [1, 2, 3])  == 3

size :: WithoutReplacement a -> Int
size (Node n _ _) = n
size (Leaf _) = 1
size Nil = 0

{-----------------------------------------------------------------------------
  Natural transformations to/from List
------------------------------------------------------------------------------}

-- | /O(n*log n)/. Build a set from a list of values.
--
-- __Duplicate values are preserved.__ Picking values from a set with
-- duplicates may return the same value as many times as that value is present
-- in the originating list.

fromList :: [a] -> WithoutReplacement a
fromList = treeFold . fmap Leaf
    where
        merge l r = let n = size l + size r in n `seq` Node n l r

        treeFold [] = Nil
        treeFold [x] = x
        treeFold xs = treeFold $ pairingWith merge xs

        pairingWith f (x:y:xs) = f x y : pairingWith f xs
        pairingWith f xs = xs

-- | /O(n*log n)/. Extract values from the set at random into a list using ST
-- to keep the state of the RNG.

toListST :: RandomGen g => WithoutReplacement a -> STGenM g s -> ST s [a]
toListST t g = do
    (res, t') <- randomWithoutReplacementM t g
    case res of
      Nothing -> return []
      Just a -> (a:) <$> toListST t' g

-- | /O(n*log n)/. Extract values from the set at random into a list using
-- the given random number generator.
--
-- >>> let gen = System.Random.mkStdGen 2022
-- >>> toList gen $ fromList [1..16]
-- [15,7,2,9,5,6,12,14,13,10,11,16,1,8,4,3]

toList :: RandomGen g => g -> WithoutReplacement a -> [a]
toList g t = runSTGen_ g $ toListST t

-- | /O(n*log n)/. Extract values from the set at random into a list using
-- the standard random number generator.
-- 
-- Because it uses the standard RNG, it must return inside a 'MonadIO'.

toListStd :: MonadIO m => WithoutReplacement a -> m [a]
toListStd t = flip toList t <$> getStdGen

{-----------------------------------------------------------------------------
  Random selection
------------------------------------------------------------------------------}

-- | /O(log n)/. Pick a random value from a set without replacement using a
-- given random number generator.
--
-- The RNG state is maintained inside a monadic environment.
--
-- The first element of the returned tuple is the picked value, which may be
-- 'Nothing' when the given set was empty.
-- The second element is the new version of the set, with the picked element
-- removed.
--
-- >>> let pureGen = System.Random.mkStdGen 2022
-- >>> gen <- System.Random.Stateful.newIOGenM pureGen
-- >>> let wr = fromList [1..16]
--
-- >>> (Just x, wr') <- randomWithoutReplacementM wr gen
-- >>> x
-- 15
--
-- >>> (Just y, wr') <- randomWithoutReplacementM wr gen
-- >>> y
-- 7

randomWithoutReplacementM :: RandomGenM g r m => WithoutReplacement a -> g -> m (Maybe a, WithoutReplacement a)
randomWithoutReplacementM Nil g = return (Nothing, Nil)
randomWithoutReplacementM (Leaf x) g = return (Just x, Nil)
randomWithoutReplacementM (Node n l r) g =
    do k <- randomRM (1, n) g
       let left = k <= size l
           b = if left then l else r
       (res, t) <- randomWithoutReplacementM b g
       let u = uncurry (Node $ n - 1) $ if left then (t, r) else (l, t)
       return (res, u)

-- | /O(log n)/. Pick a random value from a set without replacement using a
-- given random number generator.
--
-- The first element of the returned tuple is the picked value, which may be
-- 'Nothing' when the given set was empty.
-- The second element is the new version of the set, with the picked element
-- removed.
-- The third element is the updated random number generator.
--
-- >>> let gen = System.Random.mkStdGen 2022
-- >>> let wr = fromList [1..16]
--
-- >>> let (Just x, wr', gen') = randomWithoutReplacement wr gen
-- >>> x
-- 15
--
-- >>> let (Just y, wr'', gen'') = randomWithoutReplacement wr' gen'
-- >>> y
-- 7

randomWithoutReplacement :: RandomGen g => WithoutReplacement a -> g -> (Maybe a, WithoutReplacement a, g)
randomWithoutReplacement t g =
    let ((r, t'), g') = runSTGen g $ randomWithoutReplacementM t
    in (r, t', g')