{-# LANGUAGE BangPatterns #-}
-- | Optimal codes for combinatorial objects.
--
-- The integer on which a combinatorial objects is mapped is typically
-- called its rank. Below are implementations of ranking and unranking
-- algorithms for the indexes of common combinatorial objects in the
-- lexicographic enumeration of objects of the same parameters.
module Codec.Arithmetic.Combinatorics
  ( -- * Multiset Permutations

    -- | [Multiset permutations]
    -- (https://en.wikipedia.org/wiki/Permutation#Permutations_of_multisets)
    -- are ways to order the elements of a set where elements may appear
    -- more than once. The number of such permutations is equal to the
    -- multinomial coefficient with the same parameters: \[ {n \choose
    -- k_{1}, k_{2}, \ldots, k_{m}} = \frac{n!}{k_{1}! k_{2}! \cdots
    -- k_{m}!} ~~~~~\mathrm{where}~~~~~ n = \sum_i k_i \]

    rankMultisetPermutation
  , unrankMultisetPermutation
  , multinomial

  -- * Permutations

  -- | A [permutation](https://en.wikipedia.org/wiki/Permutation) is an
  -- ordering of the objects of a set of distinct elements. The number
  -- of permutations of a set of \(n\) elements is \(n!\).

  , rankPermutation
  , unrankPermutation

  -- * Combinations

  -- | A [combination](https://en.wikipedia.org/wiki/Combination) is a
  -- selection of \(k\) elements from a set of size \(n\). The number of
  -- combinations for parameters \(n\) and \(k\) is given by the
  -- binomial coefficient: \[ {n \choose k} = \frac{n!}{k! (n-k)!}  \]

  , rankCombination
  , unrankCombination
  , choose

  -- * Distributions

  -- | A distribution (usually discussed under the name [stars and
  -- bars](https://en.wikipedia.org/wiki/Stars_and_bars_(combinatorics\)))
  -- is a way to distribute \(n\) equal elements (stars) among \(k\)
  -- bins (i.e. \(k-1\) bars ).

  , rankDistribution
  , unrankDistribution

  -- * Non-Empty Distributions

  -- | The class of distributions that have at least one element per
  -- bin.

  , rankDistribution1
  , unrankDistribution1
  ) where

import Control.Exception (assert)
import Data.Maybe (fromJust)
import qualified Data.Set as S
import qualified Data.List as L
import Data.Map.Strict (Map)
import qualified Data.Map.Strict as M
import Math.Combinatorics.Exact.Factorial (factorial)

import qualified Codec.Arithmetic.Variety as V

err :: String -> a
err :: forall a. String -> a
err = String -> a
forall a. HasCallStack => String -> a
error (String -> a) -> (String -> String) -> String -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (String
"Combinatorics." String -> String -> String
forall a. [a] -> [a] -> [a]
++)

-- | Rank a multiset permutation. Returns the count of each element in
-- the set, the rank and the total number of permutations with those
-- counts (the multinomial coefficient).
rankMultisetPermutation :: Ord a => [a] -> ([(a,Int)], (Integer, Integer))
rankMultisetPermutation :: forall a. Ord a => [a] -> ([(a, Int)], (Integer, Integer))
rankMultisetPermutation [a]
msp = ( Map a Int -> [(a, Int)]
forall k a. Map k a -> [(k, a)]
M.toList Map a Int
counts
                              , (Integer
index, Integer
coef0) )
  where
    counts :: Map a Int
counts = (Map a Int -> a -> Map a Int) -> Map a Int -> [a] -> Map a Int
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
L.foldl' (\Map a Int
m a
k -> (Int -> Int -> Int) -> a -> Int -> Map a Int -> Map a Int
forall k a. Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
M.insertWith Int -> Int -> Int
forall a. Num a => a -> a -> a
(+) a
k Int
1 Map a Int
m) Map a Int
forall k a. Map k a
M.empty [a]
msp
    total0 :: Int
total0 = Map a Int -> Int
forall a. Num a => Map a a -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum Map a Int
counts
    coef0 :: Integer
coef0 = Int -> Integer
forall a. (Integral a, Bits a) => Int -> a
factorial Int
total0
            Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`div` Map a Integer -> Integer
forall a. Num a => Map a a -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product (Int -> Integer
forall a. (Integral a, Bits a) => Int -> a
factorial (Int -> Integer) -> Map a Int -> Map a Integer
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Map a Int
counts)
    index :: Integer
index = [Integer] -> Integer
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ([Integer] -> Integer) -> [Integer] -> Integer
forall a b. (a -> b) -> a -> b
$ Integer -> Integer -> Map a Int -> [a] -> [Integer]
forall a.
Ord a =>
Integer -> Integer -> Map a Int -> [a] -> [Integer]
go (Int -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
total0) Integer
coef0 Map a Int
counts [a]
msp

    go :: Ord a => Integer -> Integer -> Map a Int -> [a] -> [Integer]
    go :: forall a.
Ord a =>
Integer -> Integer -> Map a Int -> [a] -> [Integer]
go Integer
_ Integer
_ Map a Int
_ [] = []
    go Integer
total Integer
coef Map a Int
m (a
a:[a]
as) = Map a Integer -> Integer
forall a. Num a => Map a a -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum Map a Integer
lowerSubCoefs Integer -> [Integer] -> [Integer]
forall a. a -> [a] -> [a]
:
                             Integer -> Integer -> Map a Int -> [a] -> [Integer]
forall a.
Ord a =>
Integer -> Integer -> Map a Int -> [a] -> [Integer]
go Integer
total' Integer
coef' Map a Int
m' [a]
as
      where
        (Map a Int
lt,Maybe Int
eq,Map a Int
_) = a -> Map a Int -> (Map a Int, Maybe Int, Map a Int)
forall k a. Ord k => k -> Map k a -> (Map k a, Maybe a, Map k a)
M.splitLookup a
a Map a Int
m
        total' :: Integer
total' = Integer
total Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
1 -- decrement `total` by 1
        lowerSubCoefs :: Map a Integer
lowerSubCoefs = (Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`div` Integer
total) (Integer -> Integer) -> (Int -> Integer) -> Int -> Integer
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Integer
coef Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
*) (Integer -> Integer) -> (Int -> Integer) -> Int -> Integer
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Integer) -> Map a Int -> Map a Integer
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Map a Int
lt
        n :: Int
n = Maybe Int -> Int
forall a. HasCallStack => Maybe a -> a
fromJust Maybe Int
eq
        n' :: Int
n' = Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1 -- decrement `a`'s count by 1
        coef' :: Integer
coef' = (Integer
coef Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Int -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n) Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`div` Integer
total -- rm `n` factor from denom
        m' :: Map a Int
m' = (Int -> Maybe Int) -> a -> Map a Int -> Map a Int
forall k a. Ord k => (a -> Maybe a) -> k -> Map k a -> Map k a
M.update (\Int
_ -> if Int
n' Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 then Maybe Int
forall a. Maybe a
Nothing else Int -> Maybe Int
forall a. a -> Maybe a
Just Int
n')
             a
a Map a Int
m

-- | Reconstruct a multiset permutation, given the count of each element
-- in the set and a rank.
unrankMultisetPermutation :: Ord a => [(a,Int)] -> Integer -> [a]
unrankMultisetPermutation :: forall a. Ord a => [(a, Int)] -> Integer -> [a]
unrankMultisetPermutation [(a, Int)]
l Integer
i0
  | ((a, Int) -> Bool) -> [(a, Int)] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any ((Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0) (Int -> Bool) -> ((a, Int) -> Int) -> (a, Int) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a, Int) -> Int
forall a b. (a, b) -> b
snd) [(a, Int)]
l = String -> [a]
forall a. String -> a
err' String
"negative count"
  | Integer
i0 Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Integer
0 Bool -> Bool -> Bool
|| Integer
i0 Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
>= Integer
coef0 = String -> [a]
forall a. String -> a
err' (String -> [a]) -> String -> [a]
forall a b. (a -> b) -> a -> b
$ String
"out of bounds: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ (Integer, Integer) -> String
forall a. Show a => a -> String
show (Integer
i0,Integer
coef0)
  | Bool
otherwise = Integer -> Integer -> Map a Int -> Integer -> [a]
forall {t} {a} {a}.
(Integral t, Integral a, Ord a) =>
t -> t -> Map a a -> t -> [a]
go (Int -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
total0) Integer
coef0 Map a Int
counts Integer
i0
  where
    err' :: String -> c
err' = String -> c
forall a. String -> a
err (String -> c) -> (String -> String) -> String -> c
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (String
"unrankMultisetPermutation: " String -> String -> String
forall a. [a] -> [a] -> [a]
++)
    counts :: Map a Int
counts = [(a, Int)] -> Map a Int
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(a, Int)] -> Map a Int) -> [(a, Int)] -> Map a Int
forall a b. (a -> b) -> a -> b
$ ((a, Int) -> Bool) -> [(a, Int)] -> [(a, Int)]
forall a. (a -> Bool) -> [a] -> [a]
filter ((Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0) (Int -> Bool) -> ((a, Int) -> Int) -> (a, Int) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a, Int) -> Int
forall a b. (a, b) -> b
snd) [(a, Int)]
l
    total0 :: Int
total0 = Map a Int -> Int
forall a. Num a => Map a a -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum Map a Int
counts
    coef0 :: Integer
coef0 = Int -> Integer
forall a. (Integral a, Bits a) => Int -> a
factorial Int
total0
            Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`div` Map a Integer -> Integer
forall a. Num a => Map a a -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product (Int -> Integer
forall a. (Integral a, Bits a) => Int -> a
factorial (Int -> Integer) -> Map a Int -> Map a Integer
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Map a Int
counts)

    go :: t -> t -> Map a a -> t -> [a]
go t
total t
coef Map a a
m t
i | Map a a -> Bool
forall k a. Map k a -> Bool
M.null Map a a
m = []
                      | Bool
otherwise = a
a a -> [a] -> [a]
forall a. a -> [a] -> [a]
: t -> t -> Map a a -> t -> [a]
go t
total' t
coef' Map a a
m' t
i'
      where
        total' :: t
total' = t
total t -> t -> t
forall a. Num a => a -> a -> a
- t
1 -- decrement `total` by 1
        subCoefs :: Map a t
subCoefs = (t -> t -> t
forall a. Integral a => a -> a -> a
`div` t
total) (t -> t) -> (a -> t) -> a -> t
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (t
coef t -> t -> t
forall a. Num a => a -> a -> a
*) (t -> t) -> (a -> t) -> a -> t
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> t
forall a b. (Integral a, Num b) => a -> b
fromIntegral (a -> t) -> Map a a -> Map a t
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Map a a
m
        (a
a, t
lowerSubCoefsSum, t
coef') = t -> [(a, t)] -> (a, t, t)
forall {a}. t -> [(a, t)] -> (a, t, t)
findBin t
0 ([(a, t)] -> (a, t, t)) -> [(a, t)] -> (a, t, t)
forall a b. (a -> b) -> a -> b
$ Map a t -> [(a, t)]
forall k a. Map k a -> [(k, a)]
M.toList Map a t
subCoefs
        i' :: t
i' = t
i t -> t -> t
forall a. Num a => a -> a -> a
- t
lowerSubCoefsSum -- update index to local bin
        m' :: Map a a
m' = (a -> Maybe a) -> a -> Map a a -> Map a a
forall k a. Ord k => (a -> Maybe a) -> k -> Map k a -> Map k a
M.update (\a
n -> if a
n a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
1 then Maybe a
forall a. Maybe a
Nothing else a -> Maybe a
forall a. a -> Maybe a
Just (a -> Maybe a) -> a -> Maybe a
forall a b. (a -> b) -> a -> b
$ a
n a -> a -> a
forall a. Num a => a -> a -> a
- a
1)
             a
a Map a a
m

        findBin :: t -> [(a, t)] -> (a, t, t)
findBin t
_ [] = String -> (a, t, t)
forall a. String -> a
err String
"impossible"
        findBin t
acc ((a
el,t
subCoef):[(a, t)]
ascs)
          | [(a, t)] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [(a, t)]
ascs Bool -> Bool -> Bool
|| t
acc' t -> t -> Bool
forall a. Ord a => a -> a -> Bool
> t
i = (a
el, t
acc, t
subCoef)
          | Bool
otherwise = t -> [(a, t)] -> (a, t, t)
findBin t
acc' [(a, t)]
ascs
          where acc' :: t
acc' = t
acc t -> t -> t
forall a. Num a => a -> a -> a
+ t
subCoef

-- | Computes the multinomial coefficient given a list of counts \(k_i\).
multinomial :: [Int] -> Integer
multinomial :: [Int] -> Integer
multinomial [Int]
ns | (Int -> Bool) -> [Int] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0) [Int]
ns = Integer
0
               | Bool
otherwise = Int -> Integer
forall a. (Integral a, Bits a) => Int -> a
factorial ([Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Int]
ns)
                             Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`div` [Integer] -> Integer
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product (Int -> Integer
forall a. (Integral a, Bits a) => Int -> a
factorial (Int -> Integer) -> [Int] -> [Integer]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Int]
ns)

-- | Rank a permutation. Returns the rank and the total number of
-- permutations of sets with that size ( \(n!\) ).
rankPermutation :: Ord a => [a] -> (Integer, Integer)
rankPermutation :: forall a. Ord a => [a] -> (Integer, Integer)
rankPermutation [a]
p | [a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
p Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
n0 = String -> (Integer, Integer)
forall a. String -> a
err' String
"not unique elements"
                  | Bool
otherwise = Value -> (Integer, Integer)
V.fromValue Value
val
  where
    err' :: String -> c
err' = String -> c
forall a. String -> a
err (String -> c) -> (String -> String) -> String -> c
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (String
"rankPermutation: " String -> String -> String
forall a. [a] -> [a] -> [a]
++)
    s0 :: Set a
s0 = [a] -> Set a
forall a. Ord a => [a] -> Set a
S.fromList [a]
p
    n0 :: Int
n0 = Set a -> Int
forall a. Set a -> Int
S.size Set a
s0
    ns :: [Integer]
ns = Int -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Integer) -> [Int] -> [Integer]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Int
n0,Int
n0Int -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1..Int
1]
    is :: [Integer]
is = Int -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Integer) -> [Int] -> [Integer]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Set a -> [a] -> [Int]
forall {a}. Ord a => Set a -> [a] -> [Int]
go Set a
s0 [a]
p
    val :: Value
val = Bool -> ([Value] -> Value) -> [Value] -> Value
forall a. HasCallStack => Bool -> a -> a
assert ([Integer] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Integer]
is Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== [Integer] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Integer]
ns)
          [Value] -> Value
forall a. Monoid a => [a] -> a
mconcat ([Value] -> Value) -> [Value] -> Value
forall a b. (a -> b) -> a -> b
$
          (Integer -> Integer -> Value) -> [Integer] -> [Integer] -> [Value]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Integer -> Integer -> Value
V.mkValue [Integer]
is [Integer]
ns

    -- | Lookup element index in the set of remaining elements
    go :: Set a -> [a] -> [Int]
go Set a
s [] = Bool -> [Int] -> [Int]
forall a. HasCallStack => Bool -> a -> a
assert (Set a -> Bool
forall a. Set a -> Bool
S.null Set a
s) []
    go Set a
s (a
a:[a]
rest) = Int
i Int -> [Int] -> [Int]
forall a. a -> [a] -> [a]
: Set a -> [a] -> [Int]
go Set a
s' [a]
rest
      where i :: Int
i = a -> Set a -> Int
forall a. Ord a => a -> Set a -> Int
S.findIndex a
a Set a
s
            s' :: Set a
s' = a -> Set a -> Set a
forall a. Ord a => a -> Set a -> Set a
S.delete a
a Set a
s

-- | Reconstruct a permutation given a set of elements and a rank. The
-- order in which the elements of the set is given does not matter.
unrankPermutation :: Ord a => [a] -> Integer -> [a]
unrankPermutation :: forall a. Ord a => [a] -> Integer -> [a]
unrankPermutation [a]
as Integer
index
  | [a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
as Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
n = String -> [a]
forall a. String -> a
err' String
"not unique elements"
  | Integer
index Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Integer
0 Bool -> Bool -> Bool
|| Integer
index Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
>= Integer
base = String -> [a]
forall a. String -> a
err' (String -> [a]) -> String -> [a]
forall a b. (a -> b) -> a -> b
$ String
"out of bounds" String -> String -> String
forall a. [a] -> [a] -> [a]
++ (Integer, Integer) -> String
forall a. Show a => a -> String
show (Integer
index,Integer
base)
  | Bool
otherwise = Set a -> [Int] -> [a]
forall {a}. Set a -> [Int] -> [a]
go Set a
set [Int]
is
  where
    err' :: String -> c
err' = String -> c
forall a. String -> a
err (String -> c) -> (String -> String) -> String -> c
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (String
"unrankPermutation: " String -> String -> String
forall a. [a] -> [a] -> [a]
++)
    set :: Set a
set = [a] -> Set a
forall a. Ord a => [a] -> Set a
S.fromList [a]
as
    n :: Int
n = Set a -> Int
forall a. Set a -> Int
S.size Set a
set
    ns :: [Integer]
ns = Int -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Integer) -> [Int] -> [Integer]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Int
n,Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1..Int
1]
    base :: Integer
base = Int -> Integer
forall a. (Integral a, Bits a) => Int -> a
factorial (Int -> Integer) -> Int -> Integer
forall a b. (a -> b) -> a -> b
$ Int -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n
    bv :: BitVec
bv = Value -> BitVec
V.toBitVec (Value -> BitVec) -> Value -> BitVec
forall a b. (a -> b) -> a -> b
$ Integer -> Integer -> Value
V.mkValue Integer
index Integer
base
    is :: [Int]
is = Integer -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Integer -> Int) -> [Integer] -> [Int]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Integer] -> BitVec -> [Integer]
V.decode [Integer]
ns BitVec
bv

    -- | Successively delete elements at given indexes from a set
    go :: Set a -> [Int] -> [a]
go Set a
s [] = Bool -> [a] -> [a]
forall a. HasCallStack => Bool -> a -> a
assert (Set a -> Bool
forall a. Set a -> Bool
S.null Set a
s) []
    go Set a
s (Int
i:[Int]
rest) = Int -> Set a -> a
forall a. Int -> Set a -> a
S.elemAt Int
i Set a
s a -> [a] -> [a]
forall a. a -> [a] -> [a]
: Set a -> [Int] -> [a]
go (Int -> Set a -> Set a
forall a. Int -> Set a -> Set a
S.deleteAt Int
i Set a
s)  [Int]
rest

-- | Rank a combination in the form of a list of booleans. Returns the
-- \((n,k)\) parameters (where \(k\) is the number of `True` values and
-- \(n\) is the total), the rank and the total number of combinations
-- with those parameters (the binomial coefficient).
rankCombination :: [Bool] -> ((Int, Int), (Integer, Integer))
rankCombination :: [Bool] -> ((Int, Int), (Integer, Integer))
rankCombination [Bool]
c = ( (Int
n0, Int
k0)
                    , (Integer
res, Integer
n0Ck0) )
  where
    n0 :: Int
n0 = [Bool] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Bool]
c
    k0 :: Int
k0 = [Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ([Int] -> Int) -> [Int] -> Int
forall a b. (a -> b) -> a -> b
$ Bool -> Int
forall a. Enum a => a -> Int
fromEnum (Bool -> Int) -> [Bool] -> [Int]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Bool]
c
    n0Ck0 :: Integer
n0Ck0 = Int
n0 Int -> Int -> Integer
`choose` Int
k0
    res :: Integer
res = [Integer] -> Integer
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ([Integer] -> Integer) -> [Integer] -> Integer
forall a b. (a -> b) -> a -> b
$ Integer -> Integer -> Integer -> [Bool] -> [Integer]
go (Int -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n0) (Int -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
k0) Integer
n0Ck0 [Bool]
c

    go :: Integer -> Integer -> Integer -> [Bool] -> [Integer]
    go :: Integer -> Integer -> Integer -> [Bool] -> [Integer]
go Integer
_ Integer
_ Integer
_ [] = []
    go Integer
n Integer
k Integer
nCk (Bool
b:[Bool]
bs) = if Bool
b then Integer
nCk0 Integer -> [Integer] -> [Integer]
forall a. a -> [a] -> [a]
: Integer -> Integer -> Integer -> [Bool] -> [Integer]
go (Integer
nInteger -> Integer -> Integer
forall a. Num a => a -> a -> a
-Integer
1) (Integer
kInteger -> Integer -> Integer
forall a. Num a => a -> a -> a
-Integer
1) Integer
nCk1 [Bool]
bs
                        else Integer -> Integer -> Integer -> [Bool] -> [Integer]
go (Integer
nInteger -> Integer -> Integer
forall a. Num a => a -> a -> a
-Integer
1) Integer
k Integer
nCk0 [Bool]
bs
      where
        nCk0 :: Integer
nCk0 = Integer
nCk Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
nCk1 -- sub coef if 0/False
        nCk1 :: Integer
nCk1 = (Integer
nCk Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
k) Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`div` Integer
n -- sub coef if 1/True

-- | Reconstruct a combination given parameters \((n,k)\) and a rank.
unrankCombination :: (Int, Int) -> Integer -> [Bool]
unrankCombination :: (Int, Int) -> Integer -> [Bool]
unrankCombination nk :: (Int, Int)
nk@(Int
n0,Int
k0) Integer
i0
  | Int
k0 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
n0 Bool -> Bool -> Bool
|| Int
k0 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 Bool -> Bool -> Bool
|| Int
n0 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 = String -> [Bool]
forall a. String -> a
err' (String -> [Bool]) -> String -> [Bool]
forall a b. (a -> b) -> a -> b
$ String
"invalid parameters: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ (Int, Int) -> String
forall a. Show a => a -> String
show (Int, Int)
nk
  | Integer
i0 Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Integer
0 Bool -> Bool -> Bool
|| Integer
i0 Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
> Integer
n0Ck0 = String -> [Bool]
forall a. String -> a
err' (String -> [Bool]) -> String -> [Bool]
forall a b. (a -> b) -> a -> b
$ String
"out of range: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ (Integer, Integer) -> String
forall a. Show a => a -> String
show (Integer
i0,Integer
n0Ck0)
  | Bool
otherwise = Integer -> Integer -> Integer -> Integer -> [Bool]
forall {t}. Integral t => t -> t -> t -> t -> [Bool]
go (Int -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n0) (Int -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
k0) Integer
n0Ck0 Integer
i0

  where
    err' :: String -> c
err' = String -> c
forall a. String -> a
err (String -> c) -> (String -> String) -> String -> c
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (String
"unrankPermutation: " String -> String -> String
forall a. [a] -> [a] -> [a]
++)
    n0Ck0 :: Integer
n0Ck0 = Int
n0 Int -> Int -> Integer
`choose` Int
k0
    go :: t -> t -> t -> t -> [Bool]
go t
n t
k t
nCk t
i | t
n t -> t -> Bool
forall a. Eq a => a -> a -> Bool
== t
0 = []
                 | t
i t -> t -> Bool
forall a. Ord a => a -> a -> Bool
< t
nCk0 = Bool
False Bool -> [Bool] -> [Bool]
forall a. a -> [a] -> [a]
: t -> t -> t -> t -> [Bool]
go (t
nt -> t -> t
forall a. Num a => a -> a -> a
-t
1) t
k t
nCk0 t
i
                 | Bool
otherwise = Bool
True Bool -> [Bool] -> [Bool]
forall a. a -> [a] -> [a]
: t -> t -> t -> t -> [Bool]
go (t
nt -> t -> t
forall a. Num a => a -> a -> a
-t
1) (t
kt -> t -> t
forall a. Num a => a -> a -> a
-t
1) t
nCk1 (t
it -> t -> t
forall a. Num a => a -> a -> a
-t
nCk0)
      where
        nCk0 :: t
nCk0 = t
nCk t -> t -> t
forall a. Num a => a -> a -> a
- t
nCk1 -- sub coef if 0/False
        nCk1 :: t
nCk1 = (t
nCk t -> t -> t
forall a. Num a => a -> a -> a
* t
k) t -> t -> t
forall a. Integral a => a -> a -> a
`div` t
n -- sub coef if 1/True

-- | Computes the binomial coefficent given parameters \(n\) and \(k\).
choose :: Int -> Int -> Integer
choose :: Int -> Int -> Integer
choose Int
n Int
k | Integer
denom Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Integer
0 = Integer
0
           | Bool
otherwise = Integer
num Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`div` Integer
denom
  where num :: Integer
num = Int -> Integer
forall a. (Integral a, Bits a) => Int -> a
factorial Int
n
        denom :: Integer
denom = Int -> Integer
forall a. (Integral a, Bits a) => Int -> a
factorial Int
k Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Int -> Integer
forall a. (Integral a, Bits a) => Int -> a
factorial (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
k)

-- | Rank a distribution in the form of a list bin counts. Returns the
-- \((n,k)\) parameters (where \(n\) is the total number of elements and
-- \(k\) is the number of bins), the rank and the total number of
-- distributions with those parameters.
rankDistribution :: [Int] -> ((Int, Int), (Integer, Integer))
rankDistribution :: [Int] -> ((Int, Int), (Integer, Integer))
rankDistribution [] = ((Int
0,Int
0),(Integer
0,Integer
1))
rankDistribution (Int
n0:[Int]
ns)
  | Int
n0 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 Bool -> Bool -> Bool
|| (Int -> Bool) -> [Int] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0) [Int]
ns = String -> ((Int, Int), (Integer, Integer))
forall a. String -> a
err' String
"negative count"
  | Bool
otherwise = ((Int
bins,Int
balls),(Integer
i,Integer
base))
  where
    err' :: String -> c
err' = String -> c
forall a. String -> a
err (String -> c) -> (String -> String) -> String -> c
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (String
"rankDistribution: " String -> String -> String
forall a. [a] -> [a] -> [a]
++)
    comb :: [Bool]
comb = Int -> Bool -> [Bool]
forall a. Int -> a -> [a]
replicate Int
n0 Bool
False -- 0s are stars, 1s are bars
           [Bool] -> [Bool] -> [Bool]
forall a. [a] -> [a] -> [a]
++ (Int -> [Bool]) -> [Int] -> [Bool]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap ((Bool
TrueBool -> [Bool] -> [Bool]
forall a. a -> [a] -> [a]
:) ([Bool] -> [Bool]) -> (Int -> [Bool]) -> Int -> [Bool]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int -> Bool -> [Bool]) -> Bool -> Int -> [Bool]
forall a b c. (a -> b -> c) -> b -> a -> c
flip Int -> Bool -> [Bool]
forall a. Int -> a -> [a]
replicate Bool
False) [Int]
ns
    ((Int
n,Int
k),(Integer
i,Integer
base)) = [Bool] -> ((Int, Int), (Integer, Integer))
rankCombination [Bool]
comb
    bins :: Int
bins = Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
    balls :: Int
balls = Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
bins Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1

-- | Reconstruct a distribution given parameters \((n,k)\) and a rank.
unrankDistribution :: (Int, Int) -> Integer -> [Int]
unrankDistribution :: (Int, Int) -> Integer -> [Int]
unrankDistribution (Int
balls,Int
bins) Integer
i
  | Int
balls Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 Bool -> Bool -> Bool
|| Int
bins Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 = String -> [Int]
forall a. String -> a
err' (String -> [Int]) -> String -> [Int]
forall a b. (a -> b) -> a -> b
$ String
"invalid parameters: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ (Int, Int) -> String
forall a. Show a => a -> String
show (Int
balls,Int
bins)
  | Integer
i Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Integer
0 Bool -> Bool -> Bool
|| Integer
i Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
>= Integer
base = String -> [Int]
forall a. String -> a
err' (String -> [Int]) -> String -> [Int]
forall a b. (a -> b) -> a -> b
$ String
"out of range: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ (Integer, Integer) -> String
forall a. Show a => a -> String
show (Integer
i,Integer
base)
  | Int
bins Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 = []
  | Bool
otherwise = Int -> [Bool] -> [Int]
forall {t}. Num t => t -> [Bool] -> [t]
countGaps Int
0 [Bool]
bs
  where
    err' :: String -> c
err' = String -> c
forall a. String -> a
err (String -> c) -> (String -> String) -> String -> c
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (String
"unrankDistribution: " String -> String -> String
forall a. [a] -> [a] -> [a]
++)
    n :: Int
n = Int
balls Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
bins Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1 -- stars and bars
    k :: Int
k = Int
bins Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1 -- number of bars
    base :: Integer
base = if Int
bins Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 then Integer
1 else Int
n Int -> Int -> Integer
`choose` Int
k
    bs :: [Bool]
bs = (Int, Int) -> Integer -> [Bool]
unrankCombination (Int
n,Int
k) Integer
i

    countGaps :: t -> [Bool] -> [t]
countGaps !t
acc [] = [t
acc]
    countGaps !t
acc (Bool
False:[Bool]
rest) = t -> [Bool] -> [t]
countGaps (t
acc t -> t -> t
forall a. Num a => a -> a -> a
+ t
1) [Bool]
rest
    countGaps !t
acc (Bool
True:[Bool]
rest) = t
acc t -> [t] -> [t]
forall a. a -> [a] -> [a]
: t -> [Bool] -> [t]
countGaps t
0 [Bool]
rest

-- | Rank a non-empty distribution in the form of a list bin
-- counts. Returns the \((n,k)\) parameters (where \(n\) is the total
-- number of elements and \(k\) is the number of bins), the rank and the
-- total number of distributions with those parameters.
rankDistribution1 :: [Int] -> ((Int, Int), (Integer, Integer))
rankDistribution1 :: [Int] -> ((Int, Int), (Integer, Integer))
rankDistribution1 [Int]
ns
  | (Int -> Bool) -> [Int] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
1) [Int]
ns = if (Int -> Bool) -> [Int] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0) [Int]
ns then String -> ((Int, Int), (Integer, Integer))
forall a. String -> a
err' String
"negative count"
                   else String -> ((Int, Int), (Integer, Integer))
forall a. String -> a
err' String
"empty count"
  | Bool
otherwise = ((Int
balls,Int
bins),(Integer
i,Integer
base))
  where
    err' :: String -> c
err' = String -> c
forall a. String -> a
err (String -> c) -> (String -> String) -> String -> c
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (String
"rankDistribution1: " String -> String -> String
forall a. [a] -> [a] -> [a]
++)
    ((Int
balls',Int
bins),(Integer
i,Integer
base)) = [Int] -> ((Int, Int), (Integer, Integer))
rankDistribution ([Int] -> ((Int, Int), (Integer, Integer)))
-> [Int] -> ((Int, Int), (Integer, Integer))
forall a b. (a -> b) -> a -> b
$ (Int -> Int -> Int
forall a. Num a => a -> a -> a
+(-Int
1)) (Int -> Int) -> [Int] -> [Int]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Int]
ns
    balls :: Int
balls = Int
balls' Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
bins

-- | Reconstruct a distribution given parameters \((n,k)\) and a rank.
unrankDistribution1 :: (Int, Int) -> Integer -> [Int]
unrankDistribution1 :: (Int, Int) -> Integer -> [Int]
unrankDistribution1 (Int
balls,Int
bins) Integer
i
  | Int
balls Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
bins Bool -> Bool -> Bool
|| Int
bins Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 =
      String -> [Int]
forall a. String -> a
err' (String -> [Int]) -> String -> [Int]
forall a b. (a -> b) -> a -> b
$ String
"invalid parameters: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ (Int, Int) -> String
forall a. Show a => a -> String
show (Int
balls,Int
bins)
  | Bool
otherwise = (Int -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) (Int -> Int) -> [Int] -> [Int]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Int, Int) -> Integer -> [Int]
unrankDistribution (Int
balls',Int
bins) Integer
i
  where
    err' :: String -> c
err' = String -> c
forall a. String -> a
err (String -> c) -> (String -> String) -> String -> c
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (String
"unrankDistribution1: " String -> String -> String
forall a. [a] -> [a] -> [a]
++)
    balls' :: Int
balls' = Int
balls Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
bins