Skip to content

Commit

Permalink
update generation in fromListMay to match old behavior
Browse files Browse the repository at this point in the history
  • Loading branch information
byorgey committed Aug 17, 2024
1 parent 39593d1 commit d664d82
Showing 1 changed file with 11 additions and 7 deletions.
18 changes: 11 additions & 7 deletions Control/Monad/Random/Class.hs
Original file line number Diff line number Diff line change
Expand Up @@ -353,15 +353,19 @@ fromList ws = do

-- | Sample a random value from a weighted list. Return @Nothing@ if
-- the list is empty or the total weight is nonpositive.
fromListMay :: (MonadRandom m) => [(a, Rational)] -> m (Maybe a)
fromListMay :: MonadRandom m => [(a, Rational)] -> m (Maybe a)
fromListMay xs = do
let s = sum (map snd xs)
cums = scanl1 (\ ~(_,q) ~(y,s') -> (y, s'+q)) xs
case s <= 0 of
True -> return Nothing
_ -> do
let s = sum (map snd xs)
cums = scanl1 (\ ~(_, q) ~(y, s') -> (y, s' + q)) xs
if s <= 0
then return Nothing
else do
-- Pick a Word64 value uniformly
w <- getRandom
let p = s * toRational (w :: Word64) / toRational (maxBound :: Word64)
-- w / maxBound gives a uniform Rational in the range [0,1].
-- Subtract from 1 to match the way uniform Double values are
-- generated, and hence match the old behavior of this function.
let p = s * (1 - toRational (w :: Word64) / toRational (maxBound :: Word64))
return . Just . fst . head . dropWhile ((< p) . snd) $ cums

-- | Sample a value uniformly from a nonempty collection of elements.
Expand Down

0 comments on commit d664d82

Please sign in to comment.