{- |
Description: Select random elements from a list
Copyright: Copyright (C) 2023 Yoo Chung
License: GPL-3.0-or-later
Maintainer: dev@chungyc.org

Some solutions to "Problems.P23" of Ninety-Nine Haskell "Problems".
-}
module Solutions.P23 (randomSelect) where

import           Data.Set      (Set)
import qualified Data.Set      as Set
import           System.Random

-- | Extract a given number of randomly selected elements from a list.
--
-- Also return a new random number generator so that callers
-- can avoid reusing a sequence of random numbers.
randomSelect :: RandomGen g => [a] -> Int -> g -> ([a], g)
randomSelect :: forall g a. RandomGen g => [a] -> Int -> g -> ([a], g)
randomSelect [a]
xs Int
n = Set (Indexed a) -> Int -> [a] -> g -> ([a], g)
forall g a.
RandomGen g =>
Set (Indexed a) -> Int -> [a] -> g -> ([a], g)
select ([Indexed a] -> Set (Indexed a)
forall a. Ord a => [a] -> Set a
Set.fromList ([Indexed a] -> Set (Indexed a)) -> [Indexed a] -> Set (Indexed a)
forall a b. (a -> b) -> a -> b
$ (Int -> a -> Indexed a) -> [Int] -> [a] -> [Indexed a]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (((Int, a) -> Indexed a) -> Int -> a -> Indexed a
forall a b c. ((a, b) -> c) -> a -> b -> c
curry (Int, a) -> Indexed a
forall a. (Int, a) -> Indexed a
Indexed) [Int
1..] [a]
xs) Int
n []

-- | For storing an arbitrary type in 'Set'.
newtype Indexed a = Indexed (Int, a)

instance Eq (Indexed a) where
  == :: Indexed a -> Indexed a -> Bool
(==) (Indexed (Int
i, a
_)) (Indexed (Int
j, a
_)) = Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
j

instance Ord (Indexed a) where
  compare :: Indexed a -> Indexed a -> Ordering
compare (Indexed (Int
i,a
_)) (Indexed (Int
j,a
_)) = Int -> Int -> Ordering
forall a. Ord a => a -> a -> Ordering
compare Int
i Int
j

toValue :: Indexed a -> a
toValue :: forall a. Indexed a -> a
toValue (Indexed (Int
_,a
x)) = a
x

select :: RandomGen g => Set (Indexed a) -> Int -> [a] -> g -> ([a], g)
select :: forall g a.
RandomGen g =>
Set (Indexed a) -> Int -> [a] -> g -> ([a], g)
select Set (Indexed a)
_ Int
0 [a]
l g
g = ([a]
l, g
g)
select Set (Indexed a)
xs Int
n [a]
l g
g
  | Set (Indexed a) -> Bool
forall a. Set a -> Bool
Set.null Set (Indexed a)
xs = ([a]
l, g
g)
  | Bool
otherwise   = Set (Indexed a) -> Int -> [a] -> g -> ([a], g)
forall g a.
RandomGen g =>
Set (Indexed a) -> Int -> [a] -> g -> ([a], g)
select Set (Indexed a)
xs' (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) (Indexed a -> a
forall a. Indexed a -> a
toValue Indexed a
x a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a]
l) g
g'
  where (Int
i, g
g') = (Int, Int) -> g -> (Int, g)
forall g. RandomGen g => (Int, Int) -> g -> (Int, g)
forall a g. (Random a, RandomGen g) => (a, a) -> g -> (a, g)
randomR (Int
0, Set (Indexed a) -> Int
forall a. Set a -> Int
Set.size Set (Indexed a)
xs Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) g
g
        x :: Indexed a
x = Int -> Set (Indexed a) -> Indexed a
forall a. Int -> Set a -> a
Set.elemAt Int
i Set (Indexed a)
xs
        xs' :: Set (Indexed a)
xs' = Int -> Set (Indexed a) -> Set (Indexed a)
forall a. Int -> Set a -> Set a
Set.deleteAt Int
i Set (Indexed a)
xs