-- | 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
nforall a. Ord a => a -> a -> Bool
<Int
0 Bool -> Bool -> Bool
|| Int
kforall a. Ord a => a -> a -> Bool
<Int
0
  then forall a. HasCallStack => [Char] -> a
error [Char]
"choose_: n and k should nonnegative"
  else if Int
kforall a. Ord a => a -> a -> Bool
>Int
n Bool -> Bool -> Bool
|| Int
kforall a. Ord a => a -> a -> Bool
<Int
0 
    then []
    else 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 :: forall a. Int -> [a] -> [[a]]
choose Int
0 [a]
_  = [[]]
choose Int
k [] = []
choose Int
k (a
x:[a]
xs) = forall a b. (a -> b) -> [a] -> [b]
map (a
xforall a. a -> [a] -> [a]
:) (forall a. Int -> [a] -> [[a]]
choose (Int
kforall a. Num a => a -> a -> a
-Int
1) [a]
xs) forall a. [a] -> [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' :: forall a. Int -> [a] -> [([a], [a])]
choose' Int
0 [a]
xs = [([],[a]
xs)]
choose' Int
k [] = []
choose' Int
k (a
x:[a]
xs) = forall a b. (a -> b) -> [a] -> [b]
map forall {b}. ([a], b) -> ([a], b)
f (forall a. Int -> [a] -> [([a], [a])]
choose' (Int
kforall a. Num a => a -> a -> a
-Int
1) [a]
xs) forall a. [a] -> [a] -> [a]
++ forall a b. (a -> b) -> [a] -> [b]
map forall {a}. (a, [a]) -> (a, [a])
g (forall a. Int -> [a] -> [([a], [a])]
choose' Int
k [a]
xs) where
  f :: ([a], b) -> ([a], b)
f ([a]
as,b
bs) = (a
xforall a. a -> [a] -> [a]
:[a]
as ,   b
bs)
  g :: (a, [a]) -> (a, [a])
g (a
as,[a]
bs) = (  a
as , a
xforall 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'' :: forall a b. Int -> [(a, b)] -> [([a], [b])]
choose'' Int
0 [(a, b)]
xys = [([] , forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> b
snd [(a, b)]
xys)]
choose'' Int
k []  = []
choose'' Int
k ((a
x,b
y):[(a, b)]
xs) = forall a b. (a -> b) -> [a] -> [b]
map forall {b}. ([a], b) -> ([a], b)
f (forall a b. Int -> [(a, b)] -> [([a], [b])]
choose'' (Int
kforall a. Num a => a -> a -> a
-Int
1) [(a, b)]
xs) forall a. [a] -> [a] -> [a]
++ forall a b. (a -> b) -> [a] -> [b]
map forall {a}. (a, [b]) -> (a, [b])
g (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
xforall a. a -> [a] -> [a]
:[a]
as ,   b
bs)
  g :: (a, [b]) -> (a, [b])
g (a
as,[b]
bs) = (  a
as , b
yforall 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 :: forall a. Int -> [a] -> [[Either a a]]
chooseTagged Int
0 [a]
xs = [forall a b. (a -> b) -> [a] -> [b]
map forall a b. a -> Either a b
Left [a]
xs]
chooseTagged Int
k [] = []
chooseTagged Int
k (a
x:[a]
xs) = forall a b. (a -> b) -> [a] -> [b]
map forall {a}. [Either a a] -> [Either a a]
f (forall a. Int -> [a] -> [[Either a a]]
chooseTagged (Int
kforall a. Num a => a -> a -> a
-Int
1) [a]
xs) forall a. [a] -> [a] -> [a]
++ forall a b. (a -> b) -> [a] -> [b]
map forall {b}. [Either a b] -> [Either a b]
g (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 = forall a b. b -> Either a b
Right a
x forall a. a -> [a] -> [a]
: [Either a a]
eis
  g :: [Either a b] -> [Either a b]
g [Either a b]
eis = forall a b. a -> Either a b
Left  a
x 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 :: forall a. Int -> [a] -> [[a]]
combine Int
0 [a]
_  = [[]]
combine Int
k [] = []
combine Int
k xxs :: [a]
xxs@(a
x:[a]
xs) = forall a b. (a -> b) -> [a] -> [b]
map (a
xforall a. a -> [a] -> [a]
:) (forall a. Int -> [a] -> [[a]]
combine (Int
kforall a. Num a => a -> a -> a
-Int
1) [a]
xxs) forall a. [a] -> [a] -> [a]
++ forall a. Int -> [a] -> [[a]]
combine Int
k [a]
xs  

-- | A synonym for 'combine'.
compose :: Int -> [a] -> [[a]]
compose :: forall a. Int -> [a] -> [[a]]
compose = 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 :: forall a. Int -> [a] -> [[a]]
tuplesFromList Int
0 [a]
_  = [[]]
tuplesFromList Int
k [a]
xs = [ (a
yforall a. a -> [a] -> [a]
:[a]
ys) | [a]
ys <- forall a. Int -> [a] -> [[a]]
tuplesFromList (Int
kforall 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 :: forall a. [[a]] -> [[a]]
listTensor [] = [[]]
listTensor ([a]
xs:[[a]]
xss) = [ a
yforall a. a -> [a] -> [a]
:[a]
ys | [a]
ys <- 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 :: forall a. Int -> [a] -> [[a]]
kSublists = forall a. Int -> [a] -> [[a]]
choose

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

-- | All sublists of a list.
sublists :: [a] -> [[a]]
sublists :: forall a. [a] -> [[a]]
sublists [] = [[]]
sublists (a
x:[a]
xs) = forall a. [a] -> [[a]]
sublists [a]
xs forall a. [a] -> [a] -> [a]
++ forall a b. (a -> b) -> [a] -> [b]
map (a
xforall 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 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 :: forall g. RandomGen g => Int -> Int -> g -> ([Int], g)
randomChoice Int
k Int
n g
g0 = 
  if Int
kforall a. Ord a => a -> a -> Bool
>Int
n Bool -> Bool -> Bool
|| Int
kforall a. Ord a => a -> a -> Bool
<Int
0 
    then 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) = forall (t :: * -> *) s a b.
Traversable t =>
(s -> a -> (s, b)) -> s -> t a -> (s, t b)
mapAccumL (\g
g Int
m -> forall a b. (a, b) -> (b, a)
swap (forall a g. (Random a, RandomGen g) => (a, a) -> g -> (a, g)
randomR (Int
1,Int
m) g
g)) g
g0 [Int
n,Int
nforall a. Num a => a -> a -> a
-Int
1..Int
nforall a. Num a => a -> a -> a
-Int
kforall 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 a. (forall s. ST s a) -> a
runST forall a b. (a -> b) -> a -> b
$ do
    STUArray s Int Int
arr <- 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)
    forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [Int
1..Int
n] forall a b. (a -> b) -> a -> b
$ \(!Int
j) -> 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
    forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ (forall a b. [a] -> [b] -> [(a, b)]
zip [Int
n,Int
nforall a. Num a => a -> a -> a
-Int
1..] [Int]
xs) forall a b. (a -> b) -> a -> b
$ \(!Int
j,!Int
i) -> do
      Int
a <- 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 <- 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
      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
      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 <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM (forall a b. [a] -> [b] -> [(a, b)]
zip [Int
n,Int
nforall a. Num a => a -> a -> a
-Int
1..] [Int]
xs) forall a b. (a -> b) -> a -> b
$ \(!Int
j,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
    forall (m :: * -> *) a. Monad m => a -> m a
return (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 = forall a. Ord a => [a] -> [a]
sort 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' 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
xforall a. Ord a => a -> a -> Bool
>=Int
y then Int -> [Int] -> Int
skip (Int
xforall 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
xforall a. Ord a => a -> a -> Bool
<=Int
y then Int
xforall a. a -> [a] -> [a]
:Int
yforall a. a -> [a] -> [a]
:[Int]
ys else Int
y forall a. a -> [a] -> [a]
: Int -> [Int] -> [Int]
insert Int
x [Int]
ys
  insert Int
x [] = [Int
x]

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