```
-- | Subsets.

{-# LANGUAGE BangPatterns, Rank2Types #-}
module Math.Combinat.Sets
(
-- * Choices
choose_ , choose , choose' , choose'' , chooseTagged
-- * Compositions
, combine , compose
-- * Tensor products
, tuplesFromList
, listTensor
-- * Sublists
, kSublists
, sublists
, countKSublists
, countSublists
-- * Random choice
, randomChoice
)
where

--------------------------------------------------------------------------------

import Data.List ( sort , mapAccumL )
import System.Random

import Data.Array.ST
import Data.Array.MArray

-- import Data.Map (Map)
-- import qualified Data.Map as Map

import Math.Combinat.Numbers ( binomial )
import Math.Combinat.Helper  ( swap )

--------------------------------------------------------------------------------
-- * choices

-- | @choose_ k n@ returns all possible ways of choosing @k@ disjoint elements from @[1..n]@
--
-- > choose_ k n == choose k [1..n]
--
choose_ :: Int -> Int -> [[Int]]
choose_ k n  = if n<0 || k<0
then error "choose_: n and k should nonnegative"
else if k>n || k<0
then []
else choose k [1..n]

-- | All possible ways to choose @k@ elements from a list, without
-- repetitions. \"Antisymmetric power\" for lists. Synonym for 'kSublists'.
choose :: Int -> [a] -> [[a]]
choose 0 _  = [[]]
choose k [] = []
choose k (x:xs) = map (x:) (choose (k-1) xs) ++ choose k xs

-- | A version of 'choose' which also returns the complementer sets.
--
-- > choose k = map fst . choose' k
--
choose' :: Int -> [a] -> [([a],[a])]
choose' 0 xs = [([],xs)]
choose' k [] = []
choose' k (x:xs) = map f (choose' (k-1) xs) ++ map g (choose' k xs) where
f (as,bs) = (x:as ,   bs)
g (as,bs) = (  as , x:bs)

-- | Another variation of 'choose''. This satisfies
--
-- > choose'' k == map (\(xs,ys) -> (map fst xs, map snd ys)) . choose' k
--
choose'' :: Int -> [(a,b)] -> [([a],[b])]
choose'' 0 xys = [([] , map snd xys)]
choose'' k []  = []
choose'' k ((x,y):xs) = map f (choose'' (k-1) xs) ++ map g (choose'' k xs) where
f (as,bs) = (x:as ,   bs)
g (as,bs) = (  as , y:bs)

-- | Another variation on 'choose' which tags the elements based on whether they are part of
-- the selected subset ('Right') or not ('Left'):
--
-- > choose k = map rights . chooseTagged k
--
chooseTagged :: Int -> [a] -> [[Either a a]]
chooseTagged 0 xs = [map Left xs]
chooseTagged k [] = []
chooseTagged k (x:xs) = map f (chooseTagged (k-1) xs) ++ map g (chooseTagged k xs) where
f eis = Right x : eis
g eis = Left  x : eis

-- | All possible ways to choose @k@ elements from a list, /with repetitions/.
-- TODO: better name?
combine :: Int -> [a] -> [[a]]
combine 0 _  = [[]]
combine k [] = []
combine k xxs@(x:xs) = map (x:) (combine (k-1) xxs) ++ combine k xs

-- | A synonym for 'combine'.
compose :: Int -> [a] -> [[a]]
compose = combine

--------------------------------------------------------------------------------
-- * tensor products

-- | \"Tensor power\" for lists. Special case of 'listTensor':
--
-- > tuplesFromList k xs == listTensor (replicate k xs)
--
-- TODO: better name?
tuplesFromList :: Int -> [a] -> [[a]]
tuplesFromList 0 _  = [[]]
tuplesFromList k xs = [ (y:ys) | ys <- tuplesFromList (k-1) xs , y <- xs ]
--the order seems to be very important, the wrong order causes a memory leak!
--tuplesFromList k xs = [ (y:ys) | y <- xs, ys <- tuplesFromList (k-1) xs ]

-- | \"Tensor product\" for lists.
listTensor :: [[a]] -> [[a]]
listTensor [] = [[]]
listTensor (xs:xss) = [ y:ys | ys <- listTensor xss , y <- xs ]
--the order seems to be very important, the wrong order causes a memory leak!
--listTensor (xs:xss) = [ y:ys | y <- xs, ys <- listTensor xss ]

--------------------------------------------------------------------------------
-- * sublists

-- | Sublists of a list having given number of elements. Synonym for 'choose'.
kSublists :: Int -> [a] -> [[a]]
kSublists = choose

-- | @# = \binom { n } { k }@.
countKSublists :: Int -> Int -> Integer
countKSublists k n = binomial n k

-- | All sublists of a list.
sublists :: [a] -> [[a]]
sublists [] = [[]]
sublists (x:xs) = sublists xs ++ map (x:) (sublists xs)
--the order seems to be very important, the wrong order causes a memory leak!
--sublists (x:xs) = map (x:) (sublists xs) ++ sublists xs

-- | @# = 2^n@.
countSublists :: Int -> Integer
countSublists n = 2 ^ n

--------------------------------------------------------------------------------
-- * random choice

-- | @randomChoice k n@ returns a uniformly random choice of @k@ elements from the set @[1..n]@
--
-- Example:
--
-- > do
-- >   cs <- replicateM 10000 (getStdRandom (randomChoice 3 7))
-- >   mapM_ print \$ histogram cs
--
randomChoice :: RandomGen g => Int -> Int -> g -> ([Int],g)
randomChoice k n g0 =
if k>n || k<0
then error "randomChoice: k out of range"
else (makeChoiceFromIndicesKnuth n as, g1)
where
-- choose numbers from [1..n], [1..n-1], [1..n-2] etc
(g1,as) = mapAccumL (\g m -> swap (randomR (1,m) g)) g0 [n,n-1..n-k+1]

--------------------------------------------------------------------------------

-- | From a list of \$k\$ numbers, where the first is in the interval @[1..n]@,
-- the second in @[1..n-1]@, the third in @[1..n-2]@, we create a size @k@ subset of @n@.
--
-- Knuth's method. The first argument is the number @n@.
--
makeChoiceFromIndicesKnuth :: Int -> [Int] -> [Int]
makeChoiceFromIndicesKnuth n xs =
runST \$ do
arr <- newArray_ (1,n) :: forall s. ST s (STUArray s Int Int)
forM_ [1..n] \$ \(!j) -> writeArray arr j j
forM_ (zip [n,n-1..] xs) \$ \(!j,!i) -> do
writeArray arr j b
writeArray arr i a
sel <- forM (zip [n,n-1..] xs) \$ \(!j,_) -> readArray arr j
return (sort sel)

-- | From a list of \$k\$ numbers, where the first is in the interval @[1..n]@,
-- the second in @[1..n-1]@, the third in @[1..n-2]@, we create a size @k@ subset of @n@.
makeChoiceFromIndicesNaive :: [Int] -> [Int]
makeChoiceFromIndicesNaive = sort . go [] where

go :: [Int] -> [Int] -> [Int]
go acc (b:bs) = b' : go (insert b' acc) bs where b' = skip b acc
go _   [] = []

-- skip over the already occupied positions. Second argument should be a sorted list
skip :: Int -> [Int] -> Int
skip x (y:ys) = if x>=y then skip (x+1) ys else x
skip x [] = x

-- insert into a sorted list
insert :: Int -> [Int] -> [Int]
insert x (y:ys) = if x<=y then x:y:ys else y : insert x ys
insert x [] = [x]

--------------------------------------------------------------------------------

```