-- | 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 Control.Monad
import Control.Monad.ST
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_ :: Int -> Int -> [[Int]]
choose_ Int
k Int
n  = if Int
nInt -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<Int
0 Bool -> Bool -> Bool
|| Int
kInt -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<Int
0
  then [Char] -> [[Int]]
forall a. HasCallStack => [Char] -> a
error [Char]
"choose_: n and k should nonnegative"
  else if Int
kInt -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>Int
n Bool -> Bool -> Bool
|| Int
kInt -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<Int
0 
    then []
    else Int -> [Int] -> [[Int]]
forall a. Int -> [a] -> [[a]]
choose Int
k [Int
1..Int
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 :: Int -> [a] -> [[a]]
choose Int
0 [a]
_  = [[]]
choose Int
k [] = []
choose Int
k (a
x:[a]
xs) = ([a] -> [a]) -> [[a]] -> [[a]]
forall a b. (a -> b) -> [a] -> [b]
map (a
xa -> [a] -> [a]
forall a. a -> [a] -> [a]
:) (Int -> [a] -> [[a]]
forall a. Int -> [a] -> [[a]]
choose (Int
kInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) [a]
xs) [[a]] -> [[a]] -> [[a]]
forall a. [a] -> [a] -> [a]
++ Int -> [a] -> [[a]]
forall a. Int -> [a] -> [[a]]
choose Int
k [a]
xs  

-- | A version of 'choose' which also returns the complementer sets.
--
-- > choose k = map fst . choose' k
--
choose' :: Int -> [a] -> [([a],[a])]
choose' :: Int -> [a] -> [([a], [a])]
choose' Int
0 [a]
xs = [([],[a]
xs)]
choose' Int
k [] = []
choose' Int
k (a
x:[a]
xs) = (([a], [a]) -> ([a], [a])) -> [([a], [a])] -> [([a], [a])]
forall a b. (a -> b) -> [a] -> [b]
map ([a], [a]) -> ([a], [a])
forall b. ([a], b) -> ([a], b)
f (Int -> [a] -> [([a], [a])]
forall a. Int -> [a] -> [([a], [a])]
choose' (Int
kInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) [a]
xs) [([a], [a])] -> [([a], [a])] -> [([a], [a])]
forall a. [a] -> [a] -> [a]
++ (([a], [a]) -> ([a], [a])) -> [([a], [a])] -> [([a], [a])]
forall a b. (a -> b) -> [a] -> [b]
map ([a], [a]) -> ([a], [a])
forall a. (a, [a]) -> (a, [a])
g (Int -> [a] -> [([a], [a])]
forall a. Int -> [a] -> [([a], [a])]
choose' Int
k [a]
xs) where
  f :: ([a], b) -> ([a], b)
f ([a]
as,b
bs) = (a
xa -> [a] -> [a]
forall a. a -> [a] -> [a]
:[a]
as ,   b
bs)
  g :: (a, [a]) -> (a, [a])
g (a
as,[a]
bs) = (  a
as , a
xa -> [a] -> [a]
forall a. a -> [a] -> [a]
:[a]
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'' :: Int -> [(a, b)] -> [([a], [b])]
choose'' Int
0 [(a, b)]
xys = [([] , ((a, b) -> b) -> [(a, b)] -> [b]
forall a b. (a -> b) -> [a] -> [b]
map (a, b) -> b
forall a b. (a, b) -> b
snd [(a, b)]
xys)]
choose'' Int
k []  = []
choose'' Int
k ((a
x,b
y):[(a, b)]
xs) = (([a], [b]) -> ([a], [b])) -> [([a], [b])] -> [([a], [b])]
forall a b. (a -> b) -> [a] -> [b]
map ([a], [b]) -> ([a], [b])
forall b. ([a], b) -> ([a], b)
f (Int -> [(a, b)] -> [([a], [b])]
forall a b. Int -> [(a, b)] -> [([a], [b])]
choose'' (Int
kInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) [(a, b)]
xs) [([a], [b])] -> [([a], [b])] -> [([a], [b])]
forall a. [a] -> [a] -> [a]
++ (([a], [b]) -> ([a], [b])) -> [([a], [b])] -> [([a], [b])]
forall a b. (a -> b) -> [a] -> [b]
map ([a], [b]) -> ([a], [b])
forall a. (a, [b]) -> (a, [b])
g (Int -> [(a, b)] -> [([a], [b])]
forall a b. Int -> [(a, b)] -> [([a], [b])]
choose'' Int
k [(a, b)]
xs) where
  f :: ([a], b) -> ([a], b)
f ([a]
as,b
bs) = (a
xa -> [a] -> [a]
forall a. a -> [a] -> [a]
:[a]
as ,   b
bs)
  g :: (a, [b]) -> (a, [b])
g (a
as,[b]
bs) = (  a
as , b
yb -> [b] -> [b]
forall a. a -> [a] -> [a]
:[b]
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 :: Int -> [a] -> [[Either a a]]
chooseTagged Int
0 [a]
xs = [(a -> Either a a) -> [a] -> [Either a a]
forall a b. (a -> b) -> [a] -> [b]
map a -> Either a a
forall a b. a -> Either a b
Left [a]
xs]
chooseTagged Int
k [] = []
chooseTagged Int
k (a
x:[a]
xs) = ([Either a a] -> [Either a a]) -> [[Either a a]] -> [[Either a a]]
forall a b. (a -> b) -> [a] -> [b]
map [Either a a] -> [Either a a]
forall a. [Either a a] -> [Either a a]
f (Int -> [a] -> [[Either a a]]
forall a. Int -> [a] -> [[Either a a]]
chooseTagged (Int
kInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) [a]
xs) [[Either a a]] -> [[Either a a]] -> [[Either a a]]
forall a. [a] -> [a] -> [a]
++ ([Either a a] -> [Either a a]) -> [[Either a a]] -> [[Either a a]]
forall a b. (a -> b) -> [a] -> [b]
map [Either a a] -> [Either a a]
forall b. [Either a b] -> [Either a b]
g (Int -> [a] -> [[Either a a]]
forall a. Int -> [a] -> [[Either a a]]
chooseTagged Int
k [a]
xs) where
  f :: [Either a a] -> [Either a a]
f [Either a a]
eis = a -> Either a a
forall a b. b -> Either a b
Right a
x Either a a -> [Either a a] -> [Either a a]
forall a. a -> [a] -> [a]
: [Either a a]
eis
  g :: [Either a b] -> [Either a b]
g [Either a b]
eis = a -> Either a b
forall a b. a -> Either a b
Left  a
x Either a b -> [Either a b] -> [Either a b]
forall a. a -> [a] -> [a]
: [Either a b]
eis

-- | All possible ways to choose @k@ elements from a list, /with repetitions/. 
-- \"Symmetric power\" for lists. See also "Math.Combinat.Compositions".
-- TODO: better name?
combine :: Int -> [a] -> [[a]]
combine :: Int -> [a] -> [[a]]
combine Int
0 [a]
_  = [[]]
combine Int
k [] = []
combine Int
k xxs :: [a]
xxs@(a
x:[a]
xs) = ([a] -> [a]) -> [[a]] -> [[a]]
forall a b. (a -> b) -> [a] -> [b]
map (a
xa -> [a] -> [a]
forall a. a -> [a] -> [a]
:) (Int -> [a] -> [[a]]
forall a. Int -> [a] -> [[a]]
combine (Int
kInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) [a]
xxs) [[a]] -> [[a]] -> [[a]]
forall a. [a] -> [a] -> [a]
++ Int -> [a] -> [[a]]
forall a. Int -> [a] -> [[a]]
combine Int
k [a]
xs  

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

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

-- | \"Tensor power\" for lists. Special case of 'listTensor':
--
-- > tuplesFromList k xs == listTensor (replicate k xs)
-- 
-- See also "Math.Combinat.Tuples".
-- TODO: better name?
tuplesFromList :: Int -> [a] -> [[a]]
tuplesFromList :: Int -> [a] -> [[a]]
tuplesFromList Int
0 [a]
_  = [[]]
tuplesFromList Int
k [a]
xs = [ (a
ya -> [a] -> [a]
forall a. a -> [a] -> [a]
:[a]
ys) | [a]
ys <- Int -> [a] -> [[a]]
forall a. Int -> [a] -> [[a]]
tuplesFromList (Int
kInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) [a]
xs , a
y <- [a]
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 :: [[a]] -> [[a]]
listTensor [] = [[]]
listTensor ([a]
xs:[[a]]
xss) = [ a
ya -> [a] -> [a]
forall a. a -> [a] -> [a]
:[a]
ys | [a]
ys <- [[a]] -> [[a]]
forall a. [[a]] -> [[a]]
listTensor [[a]]
xss , a
y <- [a]
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 :: Int -> [a] -> [[a]]
kSublists = Int -> [a] -> [[a]]
forall a. Int -> [a] -> [[a]]
choose

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

-- | All sublists of a list.
sublists :: [a] -> [[a]]
sublists :: [a] -> [[a]]
sublists [] = [[]]
sublists (a
x:[a]
xs) = [a] -> [[a]]
forall a. [a] -> [[a]]
sublists [a]
xs [[a]] -> [[a]] -> [[a]]
forall a. [a] -> [a] -> [a]
++ ([a] -> [a]) -> [[a]] -> [[a]]
forall a b. (a -> b) -> [a] -> [b]
map (a
xa -> [a] -> [a]
forall a. a -> [a] -> [a]
:) ([a] -> [[a]]
forall a. [a] -> [[a]]
sublists [a]
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 :: Int -> Integer
countSublists Int
n = Integer
2 Integer -> Int -> Integer
forall a b. (Num a, Integral b) => a -> b -> a
^ Int
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 :: Int -> Int -> g -> ([Int], g)
randomChoice Int
k Int
n g
g0 = 
  if Int
kInt -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>Int
n Bool -> Bool -> Bool
|| Int
kInt -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<Int
0 
    then [Char] -> ([Int], g)
forall a. HasCallStack => [Char] -> a
error [Char]
"randomChoice: k out of range" 
    else (Int -> [Int] -> [Int]
makeChoiceFromIndicesKnuth Int
n [Int]
as, g
g1) 
  where
    -- choose numbers from [1..n], [1..n-1], [1..n-2] etc
    (g
g1,[Int]
as) = (g -> Int -> (g, Int)) -> g -> [Int] -> (g, [Int])
forall (t :: * -> *) a b c.
Traversable t =>
(a -> b -> (a, c)) -> a -> t b -> (a, t c)
mapAccumL (\g
g Int
m -> (Int, g) -> (g, Int)
forall a b. (a, b) -> (b, a)
swap ((Int, Int) -> g -> (Int, g)
forall a g. (Random a, RandomGen g) => (a, a) -> g -> (a, g)
randomR (Int
1,Int
m) g
g)) g
g0 [Int
n,Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1..Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
kInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
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 :: Int -> [Int] -> [Int]
makeChoiceFromIndicesKnuth Int
n [Int]
xs = 
  (forall s. ST s [Int]) -> [Int]
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s [Int]) -> [Int])
-> (forall s. ST s [Int]) -> [Int]
forall a b. (a -> b) -> a -> b
$ do
    STUArray s Int Int
arr <- (Int, Int) -> ST s (STUArray s Int Int)
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
(i, i) -> m (a i e)
newArray_ (Int
1,Int
n) :: forall s. ST s (STUArray s Int Int)
    [Int] -> (Int -> ST s ()) -> ST s ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [Int
1..Int
n] ((Int -> ST s ()) -> ST s ()) -> (Int -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \(!Int
j) -> STUArray s Int Int -> Int -> Int -> ST s ()
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> e -> m ()
writeArray STUArray s Int Int
arr Int
j Int
j
    [(Int, Int)] -> ((Int, Int) -> ST s ()) -> ST s ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ ([Int] -> [Int] -> [(Int, Int)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
n,Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1..] [Int]
xs) (((Int, Int) -> ST s ()) -> ST s ())
-> ((Int, Int) -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \(!Int
j,!Int
i) -> do
      Int
a <- STUArray s Int Int -> Int -> ST s Int
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
readArray STUArray s Int Int
arr Int
j
      Int
b <- STUArray s Int Int -> Int -> ST s Int
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
readArray STUArray s Int Int
arr Int
i
      STUArray s Int Int -> Int -> Int -> ST s ()
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> e -> m ()
writeArray STUArray s Int Int
arr Int
j Int
b
      STUArray s Int Int -> Int -> Int -> ST s ()
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> e -> m ()
writeArray STUArray s Int Int
arr Int
i Int
a
    [Int]
sel <- [(Int, Int)] -> ((Int, Int) -> ST s Int) -> ST s [Int]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM ([Int] -> [Int] -> [(Int, Int)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
n,Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1..] [Int]
xs) (((Int, Int) -> ST s Int) -> ST s [Int])
-> ((Int, Int) -> ST s Int) -> ST s [Int]
forall a b. (a -> b) -> a -> b
$ \(!Int
j,Int
_) -> STUArray s Int Int -> Int -> ST s Int
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
readArray STUArray s Int Int
arr Int
j
    [Int] -> ST s [Int]
forall (m :: * -> *) a. Monad m => a -> m a
return ([Int] -> [Int]
forall a. Ord a => [a] -> [a]
sort [Int]
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 :: [Int] -> [Int]
makeChoiceFromIndicesNaive = [Int] -> [Int]
forall a. Ord a => [a] -> [a]
sort ([Int] -> [Int]) -> ([Int] -> [Int]) -> [Int] -> [Int]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Int] -> [Int] -> [Int]
go [] where

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

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

  -- insert into a sorted list
  insert :: Int -> [Int] -> [Int]
  insert :: Int -> [Int] -> [Int]
insert Int
x (Int
y:[Int]
ys) = if Int
xInt -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<=Int
y then Int
xInt -> [Int] -> [Int]
forall a. a -> [a] -> [a]
:Int
yInt -> [Int] -> [Int]
forall a. a -> [a] -> [a]
:[Int]
ys else Int
y Int -> [Int] -> [Int]
forall a. a -> [a] -> [a]
: Int -> [Int] -> [Int]
insert Int
x [Int]
ys
  insert Int
x [] = [Int
x]

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