{-# LANGUAGE DeriveFunctor         #-}
{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE MultiParamTypeClasses #-}

-- | Efficient combinatorial algorithms over multisets, including
--   generating all permutations, partitions, subsets, cycles, and
--   other combinatorial structures based on multisets.  Note that an
--   'Eq' or 'Ord' instance on the elements is /not/ required; the
--   algorithms are careful to keep track of which things are (by
--   construction) equal to which other things, so equality testing is
--   not needed.
module Math.Combinatorics.Multiset
       ( -- * The 'Multiset' type

         Count
       , Multiset(..)
       , emptyMS, singletonMS
       , consMS, (+:)

         -- ** Conversions
       , toList
       , fromList
       , fromListEq
       , fromDistinctList
       , fromCounts
       , getCounts
       , size

         -- ** Operations
       , disjUnion
       , disjUnions

         -- * Permutations

       , permutations
       , permutationsRLE

         -- * Partitions

       , Vec
       , vPartitions
       , partitions

         -- * Submultisets

       , splits
       , kSubsets

         -- * Cycles and bracelets

       , cycles
       , bracelets
       , genFixedBracelets

         -- * Miscellaneous

       , sequenceMS

       ) where

import           Control.Arrow              (first, second, (&&&), (***))
import           Control.Monad              (forM_, when)
import           Control.Monad.Trans.Writer
import qualified Data.IntMap.Strict         as IM
import           Data.List                  (group, partition, sort)
import           Data.Maybe                 (catMaybes, fromJust)

type Count = Int

-- | A multiset is represented as a list of (element, count) pairs.
--   We maintain the invariants that the counts are always positive,
--   and no element ever appears more than once.
newtype Multiset a = MS { forall a. Multiset a -> [(a, Int)]
toCounts :: [(a, Count)] }
  deriving (Int -> Multiset a -> ShowS
[Multiset a] -> ShowS
Multiset a -> String
(Int -> Multiset a -> ShowS)
-> (Multiset a -> String)
-> ([Multiset a] -> ShowS)
-> Show (Multiset a)
forall a. Show a => Int -> Multiset a -> ShowS
forall a. Show a => [Multiset a] -> ShowS
forall a. Show a => Multiset a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Multiset a] -> ShowS
$cshowList :: forall a. Show a => [Multiset a] -> ShowS
show :: Multiset a -> String
$cshow :: forall a. Show a => Multiset a -> String
showsPrec :: Int -> Multiset a -> ShowS
$cshowsPrec :: forall a. Show a => Int -> Multiset a -> ShowS
Show, (forall a b. (a -> b) -> Multiset a -> Multiset b)
-> (forall a b. a -> Multiset b -> Multiset a) -> Functor Multiset
forall a b. a -> Multiset b -> Multiset a
forall a b. (a -> b) -> Multiset a -> Multiset b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
<$ :: forall a b. a -> Multiset b -> Multiset a
$c<$ :: forall a b. a -> Multiset b -> Multiset a
fmap :: forall a b. (a -> b) -> Multiset a -> Multiset b
$cfmap :: forall a b. (a -> b) -> Multiset a -> Multiset b
Functor)

-- | Construct a 'Multiset' from a list of (element, count) pairs.
--   Precondition: the counts must all be positive, and there must not
--   be any duplicate elements.
fromCounts :: [(a, Count)] -> Multiset a
fromCounts :: forall a. [(a, Int)] -> Multiset a
fromCounts = [(a, Int)] -> Multiset a
forall a. [(a, Int)] -> Multiset a
MS

-- | Extract just the element counts from a multiset, forgetting the
--   elements.
getCounts :: Multiset a -> [Count]
getCounts :: forall a. Multiset a -> [Int]
getCounts = ((a, Int) -> Int) -> [(a, Int)] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map (a, Int) -> Int
forall a b. (a, b) -> b
snd ([(a, Int)] -> [Int])
-> (Multiset a -> [(a, Int)]) -> Multiset a -> [Int]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Multiset a -> [(a, Int)]
forall a. Multiset a -> [(a, Int)]
toCounts

-- | Compute the total size of a multiset.
size :: Multiset a -> Int
size :: forall a. Multiset a -> Int
size = [Int] -> Int
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ([Int] -> Int) -> (Multiset a -> [Int]) -> Multiset a -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Multiset a -> [Int]
forall a. Multiset a -> [Int]
getCounts

liftMS :: ([(a, Count)] -> [(b, Count)]) -> Multiset a -> Multiset b
liftMS :: forall a b. ([(a, Int)] -> [(b, Int)]) -> Multiset a -> Multiset b
liftMS [(a, Int)] -> [(b, Int)]
f (MS [(a, Int)]
m) = [(b, Int)] -> Multiset b
forall a. [(a, Int)] -> Multiset a
MS ([(a, Int)] -> [(b, Int)]
f [(a, Int)]
m)

-- | A multiset with no values in it.
emptyMS :: Multiset a
emptyMS :: forall a. Multiset a
emptyMS = [(a, Int)] -> Multiset a
forall a. [(a, Int)] -> Multiset a
MS []

-- | Create a multiset with only a single value in it.
singletonMS :: a -> Multiset a
singletonMS :: forall a. a -> Multiset a
singletonMS a
a = [(a, Int)] -> Multiset a
forall a. [(a, Int)] -> Multiset a
MS [(a
a,Int
1)]

-- | Add an element with multiplicity to a multiset.  Precondition:
--   the new element is distinct from all elements already in the
--   multiset.
consMS :: (a, Count) -> Multiset a -> Multiset a
consMS :: forall a. (a, Int) -> Multiset a -> Multiset a
consMS e :: (a, Int)
e@(a
_,Int
c) (MS [(a, Int)]
m)
  | Int
c Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0     = [(a, Int)] -> Multiset a
forall a. [(a, Int)] -> Multiset a
MS ((a, Int)
e(a, Int) -> [(a, Int)] -> [(a, Int)]
forall a. a -> [a] -> [a]
:[(a, Int)]
m)
  | Bool
otherwise = [(a, Int)] -> Multiset a
forall a. [(a, Int)] -> Multiset a
MS [(a, Int)]
m

-- | A convenient shorthand for 'consMS'.
(+:) :: (a, Count) -> Multiset a -> Multiset a
+: :: forall a. (a, Int) -> Multiset a -> Multiset a
(+:) = (a, Int) -> Multiset a -> Multiset a
forall a. (a, Int) -> Multiset a -> Multiset a
consMS

-- | Convert a multiset to a list.
toList :: Multiset a -> [a]
toList :: forall a. Multiset a -> [a]
toList = [(a, Int)] -> [a]
forall a. [(a, Int)] -> [a]
expandCounts ([(a, Int)] -> [a])
-> (Multiset a -> [(a, Int)]) -> Multiset a -> [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Multiset a -> [(a, Int)]
forall a. Multiset a -> [(a, Int)]
toCounts

expandCounts :: [(a, Count)] -> [a]
expandCounts :: forall a. [(a, Int)] -> [a]
expandCounts = ((a, Int) -> [a]) -> [(a, Int)] -> [a]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap ((a -> Int -> [a]) -> (a, Int) -> [a]
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry ((Int -> a -> [a]) -> a -> Int -> [a]
forall a b c. (a -> b -> c) -> b -> a -> c
flip Int -> a -> [a]
forall a. Int -> a -> [a]
replicate))

-- | Efficiently convert a list to a multiset, given an 'Ord' instance
--   for the elements.  This method is provided just for convenience.
--   you can also use 'fromListEq' with only an 'Eq' instance, or
--   construct 'Multiset's directly using 'fromCounts'.
fromList :: Ord a => [a] -> Multiset a
fromList :: forall a. Ord a => [a] -> Multiset a
fromList = [(a, Int)] -> Multiset a
forall a. [(a, Int)] -> Multiset a
fromCounts ([(a, Int)] -> Multiset a)
-> ([a] -> [(a, Int)]) -> [a] -> Multiset a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([a] -> (a, Int)) -> [[a]] -> [(a, Int)]
forall a b. (a -> b) -> [a] -> [b]
map ([a] -> a
forall a. [a] -> a
head ([a] -> a) -> ([a] -> Int) -> [a] -> (a, Int)
forall (a :: * -> * -> *) b c c'.
Arrow a =>
a b c -> a b c' -> a b (c, c')
&&& [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length) ([[a]] -> [(a, Int)]) -> ([a] -> [[a]]) -> [a] -> [(a, Int)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [a] -> [[a]]
forall a. Eq a => [a] -> [[a]]
group ([a] -> [[a]]) -> ([a] -> [a]) -> [a] -> [[a]]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [a] -> [a]
forall a. Ord a => [a] -> [a]
sort

-- | Convert a list to a multiset, given an 'Eq' instance for the
--   elements.
fromListEq :: Eq a => [a] -> Multiset a
fromListEq :: forall a. Eq a => [a] -> Multiset a
fromListEq = [(a, Int)] -> Multiset a
forall a. [(a, Int)] -> Multiset a
fromCounts ([(a, Int)] -> Multiset a)
-> ([a] -> [(a, Int)]) -> [a] -> Multiset a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [a] -> [(a, Int)]
forall {a}. Eq a => [a] -> [(a, Int)]
fromListEq'
  where fromListEq' :: [a] -> [(a, Int)]
fromListEq' []     = []
        fromListEq' (a
x:[a]
xs) = (a
x, Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
xEqs) (a, Int) -> [(a, Int)] -> [(a, Int)]
forall a. a -> [a] -> [a]
: [a] -> [(a, Int)]
fromListEq' [a]
xNeqs
          where
            ([a]
xEqs, [a]
xNeqs) = (a -> Bool) -> [a] -> ([a], [a])
forall a. (a -> Bool) -> [a] -> ([a], [a])
partition (a -> a -> Bool
forall a. Eq a => a -> a -> Bool
==a
x) [a]
xs

-- | Make a multiset with one copy of each element from a list of
--   distinct elements.
fromDistinctList :: [a] -> Multiset a
fromDistinctList :: forall a. [a] -> Multiset a
fromDistinctList = [(a, Int)] -> Multiset a
forall a. [(a, Int)] -> Multiset a
fromCounts ([(a, Int)] -> Multiset a)
-> ([a] -> [(a, Int)]) -> [a] -> Multiset a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> (a, Int)) -> [a] -> [(a, Int)]
forall a b. (a -> b) -> [a] -> [b]
map (\a
x -> (a
x,Int
1))

-- | Form the disjoint union of two multisets; i.e. we assume the two
--   multisets share no elements in common.
disjUnion :: Multiset a -> Multiset a -> Multiset a
disjUnion :: forall a. Multiset a -> Multiset a -> Multiset a
disjUnion (MS [(a, Int)]
xs) (MS [(a, Int)]
ys) = [(a, Int)] -> Multiset a
forall a. [(a, Int)] -> Multiset a
MS ([(a, Int)]
xs [(a, Int)] -> [(a, Int)] -> [(a, Int)]
forall a. [a] -> [a] -> [a]
++ [(a, Int)]
ys)

-- | Form the disjoint union of a collection of multisets.  We assume
--   that the multisets all have distinct elements.
disjUnions :: [Multiset a] -> Multiset a
disjUnions :: forall a. [Multiset a] -> Multiset a
disjUnions = (Multiset a -> Multiset a -> Multiset a)
-> Multiset a -> [Multiset a] -> Multiset a
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr Multiset a -> Multiset a -> Multiset a
forall a. Multiset a -> Multiset a -> Multiset a
disjUnion ([(a, Int)] -> Multiset a
forall a. [(a, Int)] -> Multiset a
MS [])

-- | In order to generate permutations of a multiset, we need to keep
--   track of the most recently used element in the permutation being
--   built, so that we don't use it again immediately.  The
--   'RMultiset' type (for \"restricted multiset\") records this
--   information, consisting of a multiset possibly paired with an
--   element (with multiplicity) which is also part of the multiset,
--   but should not be used at the beginning of permutations.
data RMultiset a = RMS (Maybe (a, Count)) [(a,Count)]
  deriving Int -> RMultiset a -> ShowS
[RMultiset a] -> ShowS
RMultiset a -> String
(Int -> RMultiset a -> ShowS)
-> (RMultiset a -> String)
-> ([RMultiset a] -> ShowS)
-> Show (RMultiset a)
forall a. Show a => Int -> RMultiset a -> ShowS
forall a. Show a => [RMultiset a] -> ShowS
forall a. Show a => RMultiset a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [RMultiset a] -> ShowS
$cshowList :: forall a. Show a => [RMultiset a] -> ShowS
show :: RMultiset a -> String
$cshow :: forall a. Show a => RMultiset a -> String
showsPrec :: Int -> RMultiset a -> ShowS
$cshowsPrec :: forall a. Show a => Int -> RMultiset a -> ShowS
Show

-- | Convert a 'Multiset' to a 'RMultiset' (with no avoided element).
toRMS :: Multiset a -> RMultiset a
toRMS :: forall a. Multiset a -> RMultiset a
toRMS = Maybe (a, Int) -> [(a, Int)] -> RMultiset a
forall a. Maybe (a, Int) -> [(a, Int)] -> RMultiset a
RMS Maybe (a, Int)
forall a. Maybe a
Nothing ([(a, Int)] -> RMultiset a)
-> (Multiset a -> [(a, Int)]) -> Multiset a -> RMultiset a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Multiset a -> [(a, Int)]
forall a. Multiset a -> [(a, Int)]
toCounts

-- | Convert a 'RMultiset' to a 'Multiset'.
fromRMS :: RMultiset a -> Multiset a
fromRMS :: forall a. RMultiset a -> Multiset a
fromRMS (RMS Maybe (a, Int)
Nothing [(a, Int)]
m)  = [(a, Int)] -> Multiset a
forall a. [(a, Int)] -> Multiset a
MS [(a, Int)]
m
fromRMS (RMS (Just (a, Int)
e) [(a, Int)]
m) = [(a, Int)] -> Multiset a
forall a. [(a, Int)] -> Multiset a
MS ((a, Int)
e(a, Int) -> [(a, Int)] -> [(a, Int)]
forall a. a -> [a] -> [a]
:[(a, Int)]
m)

-- | List all the distinct permutations of the elements of a
--   multiset.
--
--   For example, @permutations (fromList \"abb\") ==
--   [\"abb\",\"bba\",\"bab\"]@, whereas @Data.List.permutations
--   \"abb\" == [\"abb\",\"bab\",\"bba\",\"bba\",\"bab\",\"abb\"]@.
--   This function is equivalent to, but /much/ more efficient than,
--   @nub . Data.List.permutations@, and even works when the elements
--   have no 'Eq' instance.
--
--   Note that this is a specialized version of 'permutationsRLE',
--   where each run has been expanded via 'replicate'.
permutations :: Multiset a -> [[a]]
permutations :: forall a. Multiset a -> [[a]]
permutations = ([(a, Int)] -> [a]) -> [[(a, Int)]] -> [[a]]
forall a b. (a -> b) -> [a] -> [b]
map [(a, Int)] -> [a]
forall a. [(a, Int)] -> [a]
expandCounts ([[(a, Int)]] -> [[a]])
-> (Multiset a -> [[(a, Int)]]) -> Multiset a -> [[a]]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Multiset a -> [[(a, Int)]]
forall a. Multiset a -> [[(a, Int)]]
permutationsRLE

-- | List all the distinct permutations of the elements of a multiset,
--   with each permutation run-length encoded. (Note that the
--   run-length encoding is a natural byproduct of the algorithm used,
--   not a separate postprocessing step.)
--
--   For example, @permutationsRLE [('a',1), ('b',2)] ==
--   [[('a',1),('b',2)],[('b',2),('a',1)],[('b',1),('a',1),('b',1)]]@.
--
--   (Note that although the output type is newtype-equivalent to
--   @[Multiset a]@, we don't call it that since the output may
--   violate the 'Multiset' invariant that no element should appear
--   more than once.  And indeed, morally this function does not
--   output multisets at all.)
permutationsRLE :: Multiset a -> [[(a,Count)]]
permutationsRLE :: forall a. Multiset a -> [[(a, Int)]]
permutationsRLE (MS []) = [[]]
permutationsRLE Multiset a
m       = RMultiset a -> [[(a, Int)]]
forall a. RMultiset a -> [[(a, Int)]]
permutationsRLE' (Multiset a -> RMultiset a
forall a. Multiset a -> RMultiset a
toRMS Multiset a
m)

-- | List all the (run-length encoded) distinct permutations of the
--   elements of a multiset which do not start with the element to
--   avoid (if any).
permutationsRLE' :: RMultiset a -> [[(a,Count)]]

-- If only one element is left, there's only one permutation.
permutationsRLE' :: forall a. RMultiset a -> [[(a, Int)]]
permutationsRLE' (RMS Maybe (a, Int)
Nothing [(a
x,Int
n)]) = [[(a
x,Int
n)]]

-- Otherwise, select an element+multiplicity in all possible ways, and
-- concatenate the elements to all possible permutations of the
-- remaining multiset.
permutationsRLE' RMultiset a
m  = [ (a, Int)
e (a, Int) -> [(a, Int)] -> [(a, Int)]
forall a. a -> [a] -> [a]
: [(a, Int)]
p
                      | ((a, Int)
e, RMultiset a
m') <- RMultiset a -> [((a, Int), RMultiset a)]
forall a. RMultiset a -> [((a, Int), RMultiset a)]
selectRMS RMultiset a
m
                      , [(a, Int)]
p       <- RMultiset a -> [[(a, Int)]]
forall a. RMultiset a -> [[(a, Int)]]
permutationsRLE' RMultiset a
m'
                      ]

-- | Select an element + multiplicity from a multiset in all possible
--   ways, appropriately keeping track of elements to avoid at the
--   start of permutations.
selectRMS :: RMultiset a -> [((a, Count), RMultiset a)]

-- No elements to select.
selectRMS :: forall a. RMultiset a -> [((a, Int), RMultiset a)]
selectRMS (RMS Maybe (a, Int)
_ [])            = []

-- Selecting from a multiset with n copies of x, avoiding e:
selectRMS (RMS Maybe (a, Int)
e ((a
x,Int
n) : [(a, Int)]
ms))  =

  -- If we select all n copies of x, there are no copies of x left to avoid;
  -- stick e (if it exists) back into the remaining multiset.
  ((a
x,Int
n), Maybe (a, Int) -> [(a, Int)] -> RMultiset a
forall a. Maybe (a, Int) -> [(a, Int)] -> RMultiset a
RMS Maybe (a, Int)
forall a. Maybe a
Nothing ([(a, Int)]
-> ((a, Int) -> [(a, Int)]) -> Maybe (a, Int) -> [(a, Int)]
forall b a. b -> (a -> b) -> Maybe a -> b
maybe [(a, Int)]
ms ((a, Int) -> [(a, Int)] -> [(a, Int)]
forall a. a -> [a] -> [a]
:[(a, Int)]
ms) Maybe (a, Int)
e)) ((a, Int), RMultiset a)
-> [((a, Int), RMultiset a)] -> [((a, Int), RMultiset a)]
forall a. a -> [a] -> [a]
:

  -- We can also select any number of copies of x from (n-1) down to 1; in each case,
  -- we avoid the remaining copies of x and put e back into the returned multiset.
  [ ( (a
x,Int
k), Maybe (a, Int) -> [(a, Int)] -> RMultiset a
forall a. Maybe (a, Int) -> [(a, Int)] -> RMultiset a
RMS ((a, Int) -> Maybe (a, Int)
forall a. a -> Maybe a
Just (a
x,Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
k))
                 ([(a, Int)]
-> ((a, Int) -> [(a, Int)]) -> Maybe (a, Int) -> [(a, Int)]
forall b a. b -> (a -> b) -> Maybe a -> b
maybe [(a, Int)]
ms ((a, Int) -> [(a, Int)] -> [(a, Int)]
forall a. a -> [a] -> [a]
:[(a, Int)]
ms) Maybe (a, Int)
e) )
    | Int
k <- [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
2 .. Int
1]
  ] [((a, Int), RMultiset a)]
-> [((a, Int), RMultiset a)] -> [((a, Int), RMultiset a)]
forall a. [a] -> [a] -> [a]
++

  -- Finally, we can recursively choose something other than x.
  (((a, Int), RMultiset a) -> ((a, Int), RMultiset a))
-> [((a, Int), RMultiset a)] -> [((a, Int), RMultiset a)]
forall a b. (a -> b) -> [a] -> [b]
map ((RMultiset a -> RMultiset a)
-> ((a, Int), RMultiset a) -> ((a, Int), RMultiset a)
forall (a :: * -> * -> *) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second ((a, Int) -> RMultiset a -> RMultiset a
forall a. (a, Int) -> RMultiset a -> RMultiset a
consRMS (a
x,Int
n))) (RMultiset a -> [((a, Int), RMultiset a)]
forall a. RMultiset a -> [((a, Int), RMultiset a)]
selectRMS (Maybe (a, Int) -> [(a, Int)] -> RMultiset a
forall a. Maybe (a, Int) -> [(a, Int)] -> RMultiset a
RMS Maybe (a, Int)
e [(a, Int)]
ms))

consRMS :: (a, Count) -> RMultiset a -> RMultiset a
consRMS :: forall a. (a, Int) -> RMultiset a -> RMultiset a
consRMS (a, Int)
x (RMS Maybe (a, Int)
e [(a, Int)]
m) = Maybe (a, Int) -> [(a, Int)] -> RMultiset a
forall a. Maybe (a, Int) -> [(a, Int)] -> RMultiset a
RMS Maybe (a, Int)
e ((a, Int)
x(a, Int) -> [(a, Int)] -> [(a, Int)]
forall a. a -> [a] -> [a]
:[(a, Int)]
m)


-- Some QuickCheck properties.  Of course, due to combinatorial
-- explosion these are of limited utility!
-- newtype ArbCount = ArbCount Int
--   deriving (Eq, Show, Num, Real, Enum, Ord, Integral)

-- instance Arbitrary Count where
--   arbitrary = elements (map ArbCount [1..3])

-- prop_perms_distinct :: Multiset Char ArbCount -> Bool
-- prop_perms_distinct m = length ps == length (nub ps)
--   where ps = permutations m

-- prop_perms_are_perms :: Multiset Char ArbCount -> Bool
-- prop_perms_are_perms m = all ((==l') . sort) (permutations m)
--   where l' = sort (toList m)

---------------------
-- Partitions
---------------------

-- | Element count vector.
type Vec = [Count]

-- | Componentwise comparison of count vectors.
(<|=) :: Vec -> Vec -> Bool
[Int]
xs <|= :: [Int] -> [Int] -> Bool
<|= [Int]
ys = [Bool] -> Bool
forall (t :: * -> *). Foldable t => t Bool -> Bool
and ([Bool] -> Bool) -> [Bool] -> Bool
forall a b. (a -> b) -> a -> b
$ (Int -> Int -> Bool) -> [Int] -> [Int] -> [Bool]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
(<=) [Int]
xs [Int]
ys

-- | 'vZero v' produces a zero vector of the same length as @v@.
vZero :: Vec -> Vec
vZero :: [Int] -> [Int]
vZero = (Int -> Int) -> [Int] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map (Int -> Int -> Int
forall a b. a -> b -> a
const Int
0)

-- | Test for the zero vector.
vIsZero :: Vec -> Bool
vIsZero :: [Int] -> Bool
vIsZero = (Int -> Bool) -> [Int] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
==Int
0)

-- | Do vector arithmetic componentwise.
(.+.), (.-.) :: Vec -> Vec -> Vec
.+. :: [Int] -> [Int] -> [Int]
(.+.) = (Int -> Int -> Int) -> [Int] -> [Int] -> [Int]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Int -> Int -> Int
forall a. Num a => a -> a -> a
(+)
.-. :: [Int] -> [Int] -> [Int]
(.-.) = (Int -> Int -> Int) -> [Int] -> [Int] -> [Int]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (-)

-- | Multiply a count vector by a scalar.
(*.) :: Count -> Vec -> Vec
*. :: Int -> [Int] -> [Int]
(*.) Int
n = (Int -> Int) -> [Int] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
*)

-- | 'v1 `vDiv` v2' is the largest scalar multiple of 'v2' which is
--   elementwise less than or equal to 'v1'.
vDiv :: Vec -> Vec -> Count
vDiv :: [Int] -> [Int] -> Int
vDiv [Int]
v1 [Int]
v2 = [Int] -> Int
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
minimum ([Int] -> Int) -> ([Maybe Int] -> [Int]) -> [Maybe Int] -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Maybe Int] -> [Int]
forall a. [Maybe a] -> [a]
catMaybes ([Maybe Int] -> Int) -> [Maybe Int] -> Int
forall a b. (a -> b) -> a -> b
$ (Int -> Int -> Maybe Int) -> [Int] -> [Int] -> [Maybe Int]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Int -> Int -> Maybe Int
forall {a}. Integral a => a -> a -> Maybe a
zdiv [Int]
v1 [Int]
v2
  where zdiv :: a -> a -> Maybe a
zdiv a
_ a
0 = Maybe a
forall a. Maybe a
Nothing
        zdiv a
x a
y = a -> Maybe a
forall a. a -> Maybe a
Just (a -> Maybe a) -> a -> Maybe a
forall a b. (a -> b) -> a -> b
$ a
x a -> a -> a
forall a. Integral a => a -> a -> a
`div` a
y

-- | 'vInc within v' lexicographically increments 'v' with respect to
--   'within'.  For example, @vInc [2,3,5] [1,3,4] == [1,3,5]@, and
--   @vInc [2,3,5] [1,3,5] == [2,0,0]@.
vInc :: Vec -> Vec -> Vec
vInc :: [Int] -> [Int] -> [Int]
vInc [Int]
lim [Int]
v = [Int] -> [Int]
forall a. [a] -> [a]
reverse ([Int] -> [Int] -> [Int]
forall {a}. (Num a, Ord a) => [a] -> [a] -> [a]
vInc' ([Int] -> [Int]
forall a. [a] -> [a]
reverse [Int]
lim) ([Int] -> [Int]
forall a. [a] -> [a]
reverse [Int]
v))
  where vInc' :: [a] -> [a] -> [a]
vInc' [a]
_ []          = []
        vInc' [] (a
x:[a]
xs)     = a
xa -> a -> a
forall a. Num a => a -> a -> a
+a
1 a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a]
xs
        vInc' (a
l:[a]
ls) (a
x:[a]
xs) | a
x a -> a -> Bool
forall a. Ord a => a -> a -> Bool
< a
l     = a
xa -> a -> a
forall a. Num a => a -> a -> a
+a
1 a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a]
xs
                            | Bool
otherwise = a
0 a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a] -> [a] -> [a]
vInc' [a]
ls [a]
xs

-- | Generate all vector partitions, representing each partition as a
--   multiset of vectors.
--
--   This code is a slight generalization of the code published in
--
--     Brent Yorgey. \"Generating Multiset Partitions\". In: The
--     Monad.Reader, Issue 8, September 2007.
--     <http://www.haskell.org/sitewiki/images/d/dd/TMR-Issue8.pdf>
--
--   See that article for a detailed discussion of the code and how it works.
vPartitions :: Vec -> [Multiset Vec]
vPartitions :: [Int] -> [Multiset [Int]]
vPartitions [Int]
v = [Int] -> [Int] -> [Multiset [Int]]
vPart [Int]
v ([Int] -> [Int]
vZero [Int]
v) where
  vPart :: [Int] -> [Int] -> [Multiset [Int]]
vPart [Int]
v [Int]
_ | [Int] -> Bool
vIsZero [Int]
v = [[([Int], Int)] -> Multiset [Int]
forall a. [(a, Int)] -> Multiset a
MS []]
  vPart [Int]
v [Int]
vL
    | [Int]
v [Int] -> [Int] -> Bool
forall a. Ord a => a -> a -> Bool
<= [Int]
vL   = []
    | Bool
otherwise = [([Int], Int)] -> Multiset [Int]
forall a. [(a, Int)] -> Multiset a
MS [([Int]
v,Int
1)]
                Multiset [Int] -> [Multiset [Int]] -> [Multiset [Int]]
forall a. a -> [a] -> [a]
: [ ([Int]
v',Int
k) ([Int], Int) -> Multiset [Int] -> Multiset [Int]
forall a. (a, Int) -> Multiset a -> Multiset a
+: Multiset [Int]
p' | [Int]
v' <- [Int] -> [Int] -> [Int] -> [[Int]]
withinFromTo [Int]
v ([Int] -> [Int]
vHalf [Int]
v) ([Int] -> [Int] -> [Int]
vInc [Int]
v [Int]
vL)
                                 , Int
k  <- [Int
1 .. ([Int]
v [Int] -> [Int] -> Int
`vDiv` [Int]
v')]
                                 , Multiset [Int]
p' <- [Int] -> [Int] -> [Multiset [Int]]
vPart ([Int]
v [Int] -> [Int] -> [Int]
.-. (Int
k Int -> [Int] -> [Int]
*. [Int]
v')) [Int]
v' ]

-- | 'vHalf v' computes the \"lexicographic half\" of 'v', that is,
--   the vector which is the middle element (biased towards the end)
--   in a lexicographically decreasing list of all the vectors
--   elementwise no greater than 'v'.
vHalf :: Vec -> Vec
vHalf :: [Int] -> [Int]
vHalf [] = []
vHalf (Int
x:[Int]
xs) | (Int -> Bool
forall a. Integral a => a -> Bool
even Int
x) = (Int
x Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2) Int -> [Int] -> [Int]
forall a. a -> [a] -> [a]
: [Int] -> [Int]
vHalf [Int]
xs
             | Bool
otherwise = (Int
x Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2) Int -> [Int] -> [Int]
forall a. a -> [a] -> [a]
: [Int]
xs

downFrom :: a -> [a]
downFrom a
n = [a
n,(a
na -> a -> a
forall a. Num a => a -> a -> a
-a
1)..a
0]

-- | 'within m' generates a lexicographically decreasing list of
--   vectors elementwise no greater than 'm'.
within :: Vec -> [Vec]
within :: [Int] -> [[Int]]
within = [[Int]] -> [[Int]]
forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
sequence ([[Int]] -> [[Int]]) -> ([Int] -> [[Int]]) -> [Int] -> [[Int]]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int -> [Int]) -> [Int] -> [[Int]]
forall a b. (a -> b) -> [a] -> [b]
map Int -> [Int]
forall {a}. (Num a, Enum a) => a -> [a]
downFrom

-- | Clip one vector against another.
clip :: Vec -> Vec -> Vec
clip :: [Int] -> [Int] -> [Int]
clip = (Int -> Int -> Int) -> [Int] -> [Int] -> [Int]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Int -> Int -> Int
forall a. Ord a => a -> a -> a
min

-- | 'withinFromTo m s e' efficiently generates a lexicographically
--   decreasing list of vectors which are elementwise no greater than
--   'm' and lexicographically between 's' and 'e'.
withinFromTo :: Vec -> Vec -> Vec -> [Vec]
withinFromTo :: [Int] -> [Int] -> [Int] -> [[Int]]
withinFromTo [Int]
m [Int]
s [Int]
e | Bool -> Bool
not ([Int]
s [Int] -> [Int] -> Bool
<|= [Int]
m) = [Int] -> [Int] -> [Int] -> [[Int]]
withinFromTo [Int]
m ([Int] -> [Int] -> [Int]
clip [Int]
m [Int]
s) [Int]
e
withinFromTo [Int]
m [Int]
s [Int]
e | [Int]
e [Int] -> [Int] -> Bool
forall a. Ord a => a -> a -> Bool
> [Int]
s = []
withinFromTo [Int]
m [Int]
s [Int]
e = [Int] -> [Int] -> [Int] -> Bool -> Bool -> [[Int]]
forall {a}.
(Enum a, Num a, Eq a) =>
[a] -> [a] -> [a] -> Bool -> Bool -> [[a]]
wFT [Int]
m [Int]
s [Int]
e Bool
True Bool
True
  where
    wFT :: [a] -> [a] -> [a] -> Bool -> Bool -> [[a]]
wFT [] [a]
_ [a]
_ Bool
_ Bool
_ = [[]]
    wFT (a
m:[a]
ms) (a
s:[a]
ss) (a
e:[a]
es) Bool
useS Bool
useE =
        let start :: a
start = if Bool
useS then a
s else a
m
            end :: a
end   = if Bool
useE then a
e else a
0
        in
          [a
xa -> [a] -> [a]
forall a. a -> [a] -> [a]
:[a]
xs | a
x <- [a
start,(a
starta -> a -> a
forall a. Num a => a -> a -> a
-a
1)..a
end],
                  let useS' :: Bool
useS' = Bool
useS Bool -> Bool -> Bool
&& a
xa -> a -> Bool
forall a. Eq a => a -> a -> Bool
==a
s,
                  let useE' :: Bool
useE' = Bool
useE Bool -> Bool -> Bool
&& a
xa -> a -> Bool
forall a. Eq a => a -> a -> Bool
==a
e,
                  [a]
xs <- [a] -> [a] -> [a] -> Bool -> Bool -> [[a]]
wFT [a]
ms [a]
ss [a]
es Bool
useS' Bool
useE' ]

-- | Efficiently generate all distinct multiset partitions.  Note that
--   each partition is represented as a multiset of parts (each of
--   which is a multiset) in order to properly reflect the fact that
--   some parts may occur multiple times.
partitions :: Multiset a -> [Multiset (Multiset a)]
partitions :: forall a. Multiset a -> [Multiset (Multiset a)]
partitions (MS []) = [[(Multiset a, Int)] -> Multiset (Multiset a)
forall a. [(a, Int)] -> Multiset a
MS []]
partitions (MS [(a, Int)]
m)  = ((Multiset [Int] -> Multiset (Multiset a))
-> [Multiset [Int]] -> [Multiset (Multiset a)]
forall a b. (a -> b) -> [a] -> [b]
map ((Multiset [Int] -> Multiset (Multiset a))
 -> [Multiset [Int]] -> [Multiset (Multiset a)])
-> (([Int] -> Multiset a)
    -> Multiset [Int] -> Multiset (Multiset a))
-> ([Int] -> Multiset a)
-> [Multiset [Int]]
-> [Multiset (Multiset a)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([Int] -> Multiset a) -> Multiset [Int] -> Multiset (Multiset a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap) ([a] -> [Int] -> Multiset a
forall {a}. [a] -> [Int] -> Multiset a
combine [a]
elts) ([Multiset [Int]] -> [Multiset (Multiset a)])
-> [Multiset [Int]] -> [Multiset (Multiset a)]
forall a b. (a -> b) -> a -> b
$ [Int] -> [Multiset [Int]]
vPartitions [Int]
counts
  where ([a]
elts, [Int]
counts) = [(a, Int)] -> ([a], [Int])
forall a b. [(a, b)] -> ([a], [b])
unzip [(a, Int)]
m
        combine :: [a] -> [Int] -> Multiset a
combine [a]
es [Int]
cs  = [(a, Int)] -> Multiset a
forall a. [(a, Int)] -> Multiset a
MS ([(a, Int)] -> Multiset a)
-> ([(a, Int)] -> [(a, Int)]) -> [(a, Int)] -> Multiset a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((a, Int) -> Bool) -> [(a, Int)] -> [(a, Int)]
forall a. (a -> Bool) -> [a] -> [a]
filter ((Int -> Int -> Bool
forall a. Eq 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)] -> Multiset a) -> [(a, Int)] -> Multiset a
forall a b. (a -> b) -> a -> b
$ [a] -> [Int] -> [(a, Int)]
forall a b. [a] -> [b] -> [(a, b)]
zip [a]
es [Int]
cs

-- | Generate all splittings of a multiset into two submultisets,
--   i.e. all size-two partitions.
splits :: Multiset a -> [(Multiset a, Multiset a)]
splits :: forall a. Multiset a -> [(Multiset a, Multiset a)]
splits (MS [])        = [([(a, Int)] -> Multiset a
forall a. [(a, Int)] -> Multiset a
MS [], [(a, Int)] -> Multiset a
forall a. [(a, Int)] -> Multiset a
MS [])]
splits (MS ((a
x,Int
n):[(a, Int)]
m)) =
  [Int]
-> (Int -> [(Multiset a, Multiset a)])
-> [(Multiset a, Multiset a)]
forall {a} {b}. [a] -> (a -> [b]) -> [b]
for [Int
0..Int
n] ((Int -> [(Multiset a, Multiset a)]) -> [(Multiset a, Multiset a)])
-> (Int -> [(Multiset a, Multiset a)])
-> [(Multiset a, Multiset a)]
forall a b. (a -> b) -> a -> b
$ \Int
k ->
    ((Multiset a, Multiset a) -> (Multiset a, Multiset a))
-> [(Multiset a, Multiset a)] -> [(Multiset a, Multiset a)]
forall a b. (a -> b) -> [a] -> [b]
map (a -> Int -> Multiset a -> Multiset a
forall {a}. a -> Int -> Multiset a -> Multiset a
addElt a
x Int
k (Multiset a -> Multiset a)
-> (Multiset a -> Multiset a)
-> (Multiset a, Multiset a)
-> (Multiset a, Multiset a)
forall (a :: * -> * -> *) b c b' c'.
Arrow a =>
a b c -> a b' c' -> a (b, b') (c, c')
*** a -> Int -> Multiset a -> Multiset a
forall {a}. a -> Int -> Multiset a -> Multiset a
addElt a
x (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
k)) (Multiset a -> [(Multiset a, Multiset a)]
forall a. Multiset a -> [(Multiset a, Multiset a)]
splits ([(a, Int)] -> Multiset a
forall a. [(a, Int)] -> Multiset a
MS [(a, Int)]
m))

-- | Generate all size-k submultisets.
kSubsets :: Count -> Multiset a -> [Multiset a]
kSubsets :: forall a. Int -> Multiset a -> [Multiset a]
kSubsets Int
0 Multiset a
_              = [[(a, Int)] -> Multiset a
forall a. [(a, Int)] -> Multiset a
MS []]
kSubsets Int
_ (MS [])        = []
kSubsets Int
k (MS ((a
x,Int
n):[(a, Int)]
m)) =
  [Int] -> (Int -> [Multiset a]) -> [Multiset a]
forall {a} {b}. [a] -> (a -> [b]) -> [b]
for [Int
0 .. Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
k Int
n] ((Int -> [Multiset a]) -> [Multiset a])
-> (Int -> [Multiset a]) -> [Multiset a]
forall a b. (a -> b) -> a -> b
$ \Int
j ->
    (Multiset a -> Multiset a) -> [Multiset a] -> [Multiset a]
forall a b. (a -> b) -> [a] -> [b]
map (a -> Int -> Multiset a -> Multiset a
forall {a}. a -> Int -> Multiset a -> Multiset a
addElt a
x Int
j) (Int -> Multiset a -> [Multiset a]
forall a. Int -> Multiset a -> [Multiset a]
kSubsets (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
j) ([(a, Int)] -> Multiset a
forall a. [(a, Int)] -> Multiset a
MS [(a, Int)]
m))

for :: [a] -> (a -> [b]) -> [b]
for = ((a -> [b]) -> [a] -> [b]) -> [a] -> (a -> [b]) -> [b]
forall a b c. (a -> b -> c) -> b -> a -> c
flip (a -> [b]) -> [a] -> [b]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap

addElt :: a -> Int -> Multiset a -> Multiset a
addElt a
_ Int
0 = Multiset a -> Multiset a
forall a. a -> a
id
addElt a
x Int
k = ((a
x,Int
k) (a, Int) -> Multiset a -> Multiset a
forall a. (a, Int) -> Multiset a -> Multiset a
+:)

----------------------------------------------------------------------
-- Cycles (aka Necklaces)
----------------------------------------------------------------------

-- | Generate all distinct cycles, aka necklaces, with elements taken
--   from a multiset.  See J. Sawada, \"A fast algorithm to generate
--   necklaces with fixed content\", J. Theor. Comput. Sci. 301 (2003)
--   pp. 477-489.
--
--   Given the ordering on the elements of the multiset based on their
--   position in the multiset representation (with \"smaller\"
--   elements first), in @map reverse (cycles m)@, each generated
--   cycle is lexicographically smallest among all its cyclic shifts,
--   and furthermore, the cycles occur in reverse lexicographic
--   order. (It's simply more convenient/efficient to generate the
--   cycles reversed in this way, and of course we get the same set of
--   cycles either way.)
--
--   For example, @cycles (fromList \"aabbc\") ==
--   [\"cabba\",\"bcaba\",\"cbaba\",\"bbcaa\",\"bcbaa\",\"cbbaa\"]@.
cycles :: Multiset a -> [[a]]
cycles :: forall a. Multiset a -> [[a]]
cycles (MS [])         = []   -- no such thing as an empty cycle
cycles m :: Multiset a
m@(MS ((a
x1,Int
n1):[(a, Int)]
xs))
  | Int
n1 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
1    = (Int -> Int -> Int -> [(Int, a)] -> [(Int, (a, Int))] -> [[a]]
forall a.
Int -> Int -> Int -> [(Int, a)] -> [(Int, (a, Int))] -> [[a]]
cycles' Int
n Int
2 Int
1 [(Int
0,a
x1)] ([(Int, (a, Int))] -> [(Int, (a, Int))]
forall a. [a] -> [a]
reverse ([(Int, (a, Int))] -> [(Int, (a, Int))])
-> [(Int, (a, Int))] -> [(Int, (a, Int))]
forall a b. (a -> b) -> a -> b
$ [Int] -> [(a, Int)] -> [(Int, (a, Int))]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
1..] [(a, Int)]
xs))
  | Bool
otherwise =  (Int -> Int -> Int -> [(Int, a)] -> [(Int, (a, Int))] -> [[a]]
forall a.
Int -> Int -> Int -> [(Int, a)] -> [(Int, (a, Int))] -> [[a]]
cycles' Int
n Int
2 Int
1 [(Int
0,a
x1)] ([(Int, (a, Int))] -> [(Int, (a, Int))]
forall a. [a] -> [a]
reverse ([(Int, (a, Int))] -> [(Int, (a, Int))])
-> [(Int, (a, Int))] -> [(Int, (a, Int))]
forall a b. (a -> b) -> a -> b
$ [Int] -> [(a, Int)] -> [(Int, (a, Int))]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0..] ((a
x1,Int
n1Int -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1)(a, Int) -> [(a, Int)] -> [(a, Int)]
forall a. a -> [a] -> [a]
:[(a, Int)]
xs)))
  where n :: Int
n = [Int] -> Int
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ([Int] -> Int) -> (Multiset a -> [Int]) -> Multiset a -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Multiset a -> [Int]
forall a. Multiset a -> [Int]
getCounts (Multiset a -> Int) -> Multiset a -> Int
forall a b. (a -> b) -> a -> b
$ Multiset a
m

-- | The first parameter is the length of the necklaces being
--   generated.  The second parameter @p@ is the length of the longest
--   prefix of @pre@ which is a Lyndon word, i.e. an aperiodic
--   necklace.  @pre@ is the current (reversed) prefix of the
--   necklaces being generated.
cycles' :: Int -> Int -> Int -> [(Int, a)] -> [(Int, (a,Count))] -> [[a]]
cycles' :: forall a.
Int -> Int -> Int -> [(Int, a)] -> [(Int, (a, Int))] -> [[a]]
cycles' Int
n Int
_ Int
p [(Int, a)]
pre [] | Int
n Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
p Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 = [((Int, a) -> a) -> [(Int, a)] -> [a]
forall a b. (a -> b) -> [a] -> [b]
map (Int, a) -> a
forall a b. (a, b) -> b
snd [(Int, a)]
pre]
                     | Bool
otherwise      = []

cycles' Int
n Int
t Int
p [(Int, a)]
pre [(Int, (a, Int))]
xs =
  (((Int, (a, Int)) -> Bool) -> [(Int, (a, Int))] -> [(Int, (a, Int))]
forall a. (a -> Bool) -> [a] -> [a]
takeWhile ((Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>=Int
atp) (Int -> Bool)
-> ((Int, (a, Int)) -> Int) -> (Int, (a, Int)) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int, (a, Int)) -> Int
forall a b. (a, b) -> a
fst) [(Int, (a, Int))]
xs) [(Int, (a, Int))] -> ((Int, (a, Int)) -> [[a]]) -> [[a]]
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \(Int
j, (a
xj,Int
_)) ->
    Int -> Int -> Int -> [(Int, a)] -> [(Int, (a, Int))] -> [[a]]
forall a.
Int -> Int -> Int -> [(Int, a)] -> [(Int, (a, Int))] -> [[a]]
cycles' Int
n (Int
tInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) (if Int
j Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
atp then Int
p else Int
t)
      ((Int
j,a
xj)(Int, a) -> [(Int, a)] -> [(Int, a)]
forall a. a -> [a] -> [a]
:[(Int, a)]
pre)
      (Int -> [(Int, (a, Int))] -> [(Int, (a, Int))]
forall a. Int -> [(Int, (a, Int))] -> [(Int, (a, Int))]
remove Int
j [(Int, (a, Int))]
xs)
  where atp :: Int
atp = (Int, a) -> Int
forall a b. (a, b) -> a
fst ((Int, a) -> Int) -> (Int, a) -> Int
forall a b. (a -> b) -> a -> b
$ [(Int, a)]
pre [(Int, a)] -> Int -> (Int, a)
forall a. [a] -> Int -> a
!! (Int
p Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)

remove :: Int -> [(Int, (a, Int))] -> [(Int, (a, Int))]
remove :: forall a. Int -> [(Int, (a, Int))] -> [(Int, (a, Int))]
remove Int
_ [] = []
remove Int
j (x :: (Int, (a, Int))
x@(Int
j',(a
xj,Int
nj)):[(Int, (a, Int))]
xs)
  | Int
j Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
j' Bool -> Bool -> Bool
&& Int
nj Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
1 = [(Int, (a, Int))]
xs
  | Int
j Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
j'            = (Int
j',(a
xj,Int
njInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1))(Int, (a, Int)) -> [(Int, (a, Int))] -> [(Int, (a, Int))]
forall a. a -> [a] -> [a]
:[(Int, (a, Int))]
xs
  | Bool
otherwise          = (Int, (a, Int))
x(Int, (a, Int)) -> [(Int, (a, Int))] -> [(Int, (a, Int))]
forall a. a -> [a] -> [a]
:Int -> [(Int, (a, Int))] -> [(Int, (a, Int))]
forall a. Int -> [(Int, (a, Int))] -> [(Int, (a, Int))]
remove Int
j [(Int, (a, Int))]
xs

----------------------------------------------------------------------
-- Bracelets
----------------------------------------------------------------------

-- Some utilities

--------------------------------------------------
-- Indexable and Snocable classes

class Snocable p a where
  (|>) :: p -> a -> p

-- 1-based indexing
class Indexable p where
  (!) :: p -> Int -> Int

--------------------------------------------------
-- Prenecklaces

type PreNecklace = [Int]

-- A prenecklace, stored backwards, along with its length and its
-- first element cached for quick retrieval.
data Pre = Pre !Int (Maybe Int) PreNecklace
  deriving (Int -> Pre -> ShowS
[Pre] -> ShowS
Pre -> String
(Int -> Pre -> ShowS)
-> (Pre -> String) -> ([Pre] -> ShowS) -> Show Pre
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Pre] -> ShowS
$cshowList :: [Pre] -> ShowS
show :: Pre -> String
$cshow :: Pre -> String
showsPrec :: Int -> Pre -> ShowS
$cshowsPrec :: Int -> Pre -> ShowS
Show)

emptyPre :: Pre
emptyPre :: Pre
emptyPre = Int -> Maybe Int -> [Int] -> Pre
Pre Int
0 Maybe Int
forall a. Maybe a
Nothing []

getPre :: Pre -> PreNecklace
getPre :: Pre -> [Int]
getPre (Pre Int
_ Maybe Int
_ [Int]
as) = [Int] -> [Int]
forall a. [a] -> [a]
reverse [Int]
as

instance Snocable Pre Int where
  (Pre Int
0 Maybe Int
_ []) |> :: Pre -> Int -> Pre
|> Int
a  = Int -> Maybe Int -> [Int] -> Pre
Pre Int
1 (Int -> Maybe Int
forall a. a -> Maybe a
Just Int
a) [Int
a]
  (Pre Int
t Maybe Int
a1 [Int]
as) |> Int
a = Int -> Maybe Int -> [Int] -> Pre
Pre (Int
tInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) Maybe Int
a1 (Int
aInt -> [Int] -> [Int]
forall a. a -> [a] -> [a]
:[Int]
as)

instance Indexable Pre where
  Pre
_ ! :: Pre -> Int -> Int
! Int
0 = Int
0
  (Pre Int
_ (Just Int
a1) [Int]
_) ! Int
1 = Int
a1
  (Pre Int
t Maybe Int
_ [Int]
as) ! Int
i = [Int]
as [Int] -> Int -> Int
forall a. [a] -> Int -> a
!! (Int
tInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
i)
    -- as stores  a_t .. a_1.
    -- a_1 is the last element, i.e. with index t-1.
    -- a_2 has index t-2.
    -- In general, a_i has index t-i.

--------------------------------------------------
-- Run-length encoding

-- Run-length encodings.  Stored in *reverse* order for easy access to
-- the end.
data RLE a = RLE !Int !Int [(a,Int)]
  deriving (Int -> RLE a -> ShowS
[RLE a] -> ShowS
RLE a -> String
(Int -> RLE a -> ShowS)
-> (RLE a -> String) -> ([RLE a] -> ShowS) -> Show (RLE a)
forall a. Show a => Int -> RLE a -> ShowS
forall a. Show a => [RLE a] -> ShowS
forall a. Show a => RLE a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [RLE a] -> ShowS
$cshowList :: forall a. Show a => [RLE a] -> ShowS
show :: RLE a -> String
$cshow :: forall a. Show a => RLE a -> String
showsPrec :: Int -> RLE a -> ShowS
$cshowsPrec :: forall a. Show a => Int -> RLE a -> ShowS
Show)
  -- First Int is the total length of the decoded list.
  -- Second Int is the number of blocks.

emptyRLE :: RLE a
emptyRLE :: forall a. RLE a
emptyRLE = Int -> Int -> [(a, Int)] -> RLE a
forall a. Int -> Int -> [(a, Int)] -> RLE a
RLE Int
0 Int
0 []

compareRLE :: Ord a => [(a,Int)] -> [(a,Int)] -> Ordering
compareRLE :: forall a. Ord a => [(a, Int)] -> [(a, Int)] -> Ordering
compareRLE [] [] = Ordering
EQ
compareRLE [] [(a, Int)]
_  = Ordering
LT
compareRLE [(a, Int)]
_ []  = Ordering
GT
compareRLE ((a
a1,Int
n1):[(a, Int)]
rle1) ((a
a2,Int
n2):[(a, Int)]
rle2)
  | (a
a1,Int
n1) (a, Int) -> (a, Int) -> Bool
forall a. Eq a => a -> a -> Bool
== (a
a2,Int
n2) = [(a, Int)] -> [(a, Int)] -> Ordering
forall a. Ord a => [(a, Int)] -> [(a, Int)] -> Ordering
compareRLE [(a, Int)]
rle1 [(a, Int)]
rle2
  | a
a1 a -> a -> Bool
forall a. Ord a => a -> a -> Bool
< a
a2 = Ordering
LT
  | a
a1 a -> a -> Bool
forall a. Ord a => a -> a -> Bool
> a
a2 = Ordering
GT
  | (Int
n1 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
n2 Bool -> Bool -> Bool
&& ([(a, Int)] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [(a, Int)]
rle1 Bool -> Bool -> Bool
|| (a, Int) -> a
forall a b. (a, b) -> a
fst ([(a, Int)] -> (a, Int)
forall a. [a] -> a
head [(a, Int)]
rle1) a -> a -> Bool
forall a. Ord a => a -> a -> Bool
< a
a2)) Bool -> Bool -> Bool
|| (Int
n1 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
n2 Bool -> Bool -> Bool
&& Bool -> Bool
not ([(a, Int)] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [(a, Int)]
rle2) Bool -> Bool -> Bool
&& a
a1 a -> a -> Bool
forall a. Ord a => a -> a -> Bool
< (a, Int) -> a
forall a b. (a, b) -> a
fst ([(a, Int)] -> (a, Int)
forall a. [a] -> a
head [(a, Int)]
rle2)) = Ordering
LT
  | Bool
otherwise = Ordering
GT

instance Indexable (RLE Int) where
  (RLE Int
_ Int
_ []) ! :: RLE Int -> Int -> Int
! Int
_ = String -> Int
forall a. HasCallStack => String -> a
error String
"Bad index in (!) for RLE"
  (RLE Int
n Int
b ((Int
a,Int
v):[(Int, Int)]
rest)) ! Int
i
    | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
v = Int
a
    | Bool
otherwise = (Int -> Int -> [(Int, Int)] -> RLE Int
forall a. Int -> Int -> [(a, Int)] -> RLE a
RLE (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
v) (Int
bInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) [(Int, Int)]
rest) RLE Int -> Int -> Int
forall p. Indexable p => p -> Int -> Int
! (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
v)

instance Eq a => Snocable (RLE a) a where
  (RLE Int
_ Int
_ []) |> :: RLE a -> a -> RLE a
|> a
a' = Int -> Int -> [(a, Int)] -> RLE a
forall a. Int -> Int -> [(a, Int)] -> RLE a
RLE Int
1 Int
1 [(a
a',Int
1)]
  (RLE Int
n Int
b rle :: [(a, Int)]
rle@((a
a,Int
v):[(a, Int)]
rest)) |> a
a'
    | a
a a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
a'   = Int -> Int -> [(a, Int)] -> RLE a
forall a. Int -> Int -> [(a, Int)] -> RLE a
RLE (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) Int
b     ((a
a,Int
vInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1)(a, Int) -> [(a, Int)] -> [(a, Int)]
forall a. a -> [a] -> [a]
:[(a, Int)]
rest)
    | Bool
otherwise = Int -> Int -> [(a, Int)] -> RLE a
forall a. Int -> Int -> [(a, Int)] -> RLE a
RLE (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) (Int
bInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) ((a
a',Int
1)(a, Int) -> [(a, Int)] -> [(a, Int)]
forall a. a -> [a] -> [a]
:[(a, Int)]
rle)

--------------------------------------------------
-- Prenecklaces + RLE

-- Prenecklaces along with a run-length encoding.
data Pre' = Pre' Pre (RLE Int)
  deriving Int -> Pre' -> ShowS
[Pre'] -> ShowS
Pre' -> String
(Int -> Pre' -> ShowS)
-> (Pre' -> String) -> ([Pre'] -> ShowS) -> Show Pre'
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Pre'] -> ShowS
$cshowList :: [Pre'] -> ShowS
show :: Pre' -> String
$cshow :: Pre' -> String
showsPrec :: Int -> Pre' -> ShowS
$cshowsPrec :: Int -> Pre' -> ShowS
Show

emptyPre' :: Pre'
emptyPre' :: Pre'
emptyPre' = Pre -> RLE Int -> Pre'
Pre' Pre
emptyPre RLE Int
forall a. RLE a
emptyRLE

getPre' :: Pre' -> PreNecklace
getPre' :: Pre' -> [Int]
getPre' (Pre' Pre
pre RLE Int
_) = Pre -> [Int]
getPre Pre
pre

instance Indexable Pre' where
  Pre'
_ ! :: Pre' -> Int -> Int
! Int
0 = Int
0
  (Pre' (Pre Int
len Maybe Int
_ [Int]
_) RLE Int
rle) ! Int
i = RLE Int
rle RLE Int -> Int -> Int
forall p. Indexable p => p -> Int -> Int
! (Int
len Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)

instance Snocable Pre' Int where
  (Pre' Pre
p RLE Int
rle) |> :: Pre' -> Int -> Pre'
|> Int
a = Pre -> RLE Int -> Pre'
Pre' (Pre
p Pre -> Int -> Pre
forall p a. Snocable p a => p -> a -> p
|> Int
a) (RLE Int
rle RLE Int -> Int -> RLE Int
forall p a. Snocable p a => p -> a -> p
|> Int
a)

--------------------------------------------------
-- Bracelet generation

type Bracelet = [Int]

-- | An optimized bracelet generation algorithm, based on
--   S. Karim et al, "Generating Bracelets with Fixed Content".
--   <http://www.cis.uoguelph.ca/~sawada/papers/fix-brace.pdf>
--
--   @genFixedBracelets n content@ produces all bracelets (unique up
--   to rotation and reflection) of length @n@ using @content@, which
--   consists of a list of pairs where the pair (a,i) indicates that
--   element a may be used up to i times.  It is assumed that the elements
--   are drawn from [0..k].
genFixedBracelets :: Int -> [(Int,Int)] -> [Bracelet]
genFixedBracelets :: Int -> [(Int, Int)] -> [[Int]]
genFixedBracelets Int
n [(Int
0,Int
k)] | Int
k Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
n = [Int -> Int -> [Int]
forall a. Int -> a -> [a]
replicate Int
k Int
0]
                            | Bool
otherwise = []
genFixedBracelets Int
n [(Int, Int)]
content = Writer [[Int]] () -> [[Int]]
forall w a. Writer w a -> w
execWriter (Int -> Int -> Int -> IntMap Int -> Pre' -> Writer [[Int]] ()
go Int
1 Int
1 Int
0 ([(Int, Int)] -> IntMap Int
forall a. [(Int, a)] -> IntMap a
IM.fromList [(Int, Int)]
content) Pre'
emptyPre')
  where
    go :: Int -> Int -> Int -> IM.IntMap Int -> Pre' -> Writer [Bracelet] ()
    go :: Int -> Int -> Int -> IntMap Int -> Pre' -> Writer [[Int]] ()
go Int
_ Int
_ Int
_ IntMap Int
con Pre'
_ | IntMap Int -> [Int]
forall a. IntMap a -> [Int]
IM.keys IntMap Int
con [Int] -> [Int] -> Bool
forall a. Eq a => a -> a -> Bool
== [Int
0] = () -> Writer [[Int]] ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
    go Int
t Int
p Int
r IntMap Int
con pre :: Pre'
pre@(Pre' (Pre Int
_ Maybe Int
_ [Int]
as) RLE Int
_)
      | Int
t Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
n =
          Bool -> Writer [[Int]] () -> Writer [[Int]] ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int -> [Int] -> [Int]
forall a. Int -> [a] -> [a]
take (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
r) [Int]
as [Int] -> [Int] -> Bool
forall a. Ord a => a -> a -> Bool
>= [Int] -> [Int]
forall a. [a] -> [a]
reverse (Int -> [Int] -> [Int]
forall a. Int -> [a] -> [a]
take (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
r) [Int]
as) Bool -> Bool -> Bool
&& Int
n Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
p Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0) (Writer [[Int]] () -> Writer [[Int]] ())
-> Writer [[Int]] () -> Writer [[Int]] ()
forall a b. (a -> b) -> a -> b
$
            [[Int]] -> Writer [[Int]] ()
forall (m :: * -> *) w. Monad m => w -> WriterT w m ()
tell [Pre' -> [Int]
getPre' Pre'
pre]
      | Bool
otherwise = do
          let a' :: Int
a' = Pre'
pre Pre' -> Int -> Int
forall p. Indexable p => p -> Int -> Int
! (Int
tInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
p)
          [Int] -> (Int -> Writer [[Int]] ()) -> Writer [[Int]] ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ ((Int -> Bool) -> [Int] -> [Int]
forall a. (a -> Bool) -> [a] -> [a]
dropWhile (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
a') ([Int] -> [Int]) -> [Int] -> [Int]
forall a b. (a -> b) -> a -> b
$ IntMap Int -> [Int]
forall a. IntMap a -> [Int]
IM.keys IntMap Int
con) ((Int -> Writer [[Int]] ()) -> Writer [[Int]] ())
-> (Int -> Writer [[Int]] ()) -> Writer [[Int]] ()
forall a b. (a -> b) -> a -> b
$ \Int
j -> do
            let con' :: IntMap Int
con' = Int -> IntMap Int -> IntMap Int
decrease Int
j IntMap Int
con
                pre' :: Pre'
pre' = Pre'
pre Pre' -> Int -> Pre'
forall p a. Snocable p a => p -> a -> p
|> Int
j
                c :: Ordering
c = Int -> Pre' -> Ordering
forall {p}. p -> Pre' -> Ordering
checkRev2 Int
t Pre'
pre'
                p' :: Int
p' | Int
j Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
a'   = Int
t
                   | Bool
otherwise = Int
p
            Bool -> Writer [[Int]] () -> Writer [[Int]] ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Ordering
c Ordering -> Ordering -> Bool
forall a. Eq a => a -> a -> Bool
== Ordering
EQ) (Writer [[Int]] () -> Writer [[Int]] ())
-> Writer [[Int]] () -> Writer [[Int]] ()
forall a b. (a -> b) -> a -> b
$ Int -> Int -> Int -> IntMap Int -> Pre' -> Writer [[Int]] ()
go (Int
tInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) Int
p' Int
t IntMap Int
con' Pre'
pre'
            Bool -> Writer [[Int]] () -> Writer [[Int]] ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Ordering
c Ordering -> Ordering -> Bool
forall a. Eq a => a -> a -> Bool
== Ordering
GT) (Writer [[Int]] () -> Writer [[Int]] ())
-> Writer [[Int]] () -> Writer [[Int]] ()
forall a b. (a -> b) -> a -> b
$ Int -> Int -> Int -> IntMap Int -> Pre' -> Writer [[Int]] ()
go (Int
tInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) Int
p' Int
r IntMap Int
con' Pre'
pre'

    decrease :: Int -> IM.IntMap Int -> IM.IntMap Int
    decrease :: Int -> IntMap Int -> IntMap Int
decrease Int
j IntMap Int
con
      | IntMap Int -> Bool
forall a. IntMap a -> Bool
IM.null IntMap Int
con = IntMap Int
con
      | Bool
otherwise   = (Maybe Int -> Maybe Int) -> Int -> IntMap Int -> IntMap Int
forall a. (Maybe a -> Maybe a) -> Int -> IntMap a -> IntMap a
IM.alter Maybe Int -> Maybe Int
forall {a}. (Eq a, Num a) => Maybe a -> Maybe a
q Int
j IntMap Int
con
      where
        q :: Maybe a -> Maybe a
q (Just a
1)   = Maybe a
forall a. Maybe a
Nothing
        q (Just a
cnt) = a -> Maybe a
forall a. a -> Maybe a
Just (a
cnta -> a -> a
forall a. Num a => a -> a -> a
-a
1)
        q Maybe a
_          = Maybe a
forall a. Maybe a
Nothing

    checkRev2 :: p -> Pre' -> Ordering
checkRev2 p
_ (Pre' Pre
_ (RLE Int
_ Int
_ [(Int, Int)]
rle)) = [(Int, Int)] -> [(Int, Int)] -> Ordering
forall a. Ord a => [(a, Int)] -> [(a, Int)] -> Ordering
compareRLE [(Int, Int)]
rle ([(Int, Int)] -> [(Int, Int)]
forall a. [a] -> [a]
reverse [(Int, Int)]
rle)

-- | Generate all distinct bracelets (lists considered equivalent up
--   to rotation and reversal) from a given multiset.  The generated
--   bracelets are in lexicographic order, and each is
--   lexicographically smallest among its rotations and reversals.
--   See @genFixedBracelets@ for a slightly more general routine with
--   references.
--
--   For example, @bracelets $ fromList \"RRRRRRRLLL\"@ yields
--
--   > ["LLLRRRRRRR","LLRLRRRRRR","LLRRLRRRRR","LLRRRLRRRR"
--   > ,"LRLRLRRRRR","LRLRRLRRRR","LRLRRRLRRR","LRRLRRLRRR"]
bracelets :: Multiset a -> [[a]]
bracelets :: forall a. Multiset a -> [[a]]
bracelets ms :: Multiset a
ms@(MS [(a, Int)]
cnts) = [[a]]
bs
  where
    contentMap :: IntMap a
contentMap = [(Int, a)] -> IntMap a
forall a. [(Int, a)] -> IntMap a
IM.fromList ([Int] -> [a] -> [(Int, a)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0..] (((a, Int) -> a) -> [(a, Int)] -> [a]
forall a b. (a -> b) -> [a] -> [b]
map (a, Int) -> a
forall a b. (a, b) -> a
fst [(a, Int)]
cnts))
    content :: [(Int, Int)]
content    = (Int -> (a, Int) -> (Int, Int))
-> [Int] -> [(a, Int)] -> [(Int, Int)]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\Int
i (a
_,Int
n) -> (Int
i,Int
n)) [Int
0..] [(a, Int)]
cnts
    rawBs :: [[Int]]
rawBs      = Int -> [(Int, Int)] -> [[Int]]
genFixedBracelets (Multiset a -> Int
forall a. Multiset a -> Int
size Multiset a
ms) [(Int, Int)]
content
    bs :: [[a]]
bs         = ([Int] -> [a]) -> [[Int]] -> [[a]]
forall a b. (a -> b) -> [a] -> [b]
map ((Int -> a) -> [Int] -> [a]
forall a b. (a -> b) -> [a] -> [b]
map (Maybe a -> a
forall a. HasCallStack => Maybe a -> a
fromJust (Maybe a -> a) -> (Int -> Maybe a) -> Int -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int -> IntMap a -> Maybe a) -> IntMap a -> Int -> Maybe a
forall a b c. (a -> b -> c) -> b -> a -> c
flip Int -> IntMap a -> Maybe a
forall a. Int -> IntMap a -> Maybe a
IM.lookup IntMap a
contentMap)) [[Int]]
rawBs

----------------------------------------------------------------------
-- sequenceMS
----------------------------------------------------------------------

-- | Take a multiset of lists, and select one element from each list
--   in every possible combination to form a list of multisets.  We
--   assume that all the list elements are distinct.
sequenceMS :: Multiset [a] -> [Multiset a]
sequenceMS :: forall a. Multiset [a] -> [Multiset a]
sequenceMS = ([Multiset a] -> Multiset a) -> [[Multiset a]] -> [Multiset a]
forall a b. (a -> b) -> [a] -> [b]
map [Multiset a] -> Multiset a
forall a. [Multiset a] -> Multiset a
disjUnions
           ([[Multiset a]] -> [Multiset a])
-> (Multiset [a] -> [[Multiset a]]) -> Multiset [a] -> [Multiset a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [[Multiset a]] -> [[Multiset a]]
forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
sequence
           ([[Multiset a]] -> [[Multiset a]])
-> (Multiset [a] -> [[Multiset a]])
-> Multiset [a]
-> [[Multiset a]]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (([a], Int) -> [Multiset a]) -> [([a], Int)] -> [[Multiset a]]
forall a b. (a -> b) -> [a] -> [b]
map (\([a]
xs, Int
n) -> Int -> Multiset a -> [Multiset a]
forall a. Int -> Multiset a -> [Multiset a]
kSubsets Int
n ([(a, Int)] -> Multiset a
forall a. [(a, Int)] -> Multiset a
MS ([(a, Int)] -> Multiset a) -> [(a, Int)] -> Multiset a
forall a b. (a -> b) -> a -> b
$ ([a], Int) -> [(a, Int)]
forall a. ([a], Int) -> [(a, Int)]
uncollate ([a]
xs, Int
n)))
           ([([a], Int)] -> [[Multiset a]])
-> (Multiset [a] -> [([a], Int)]) -> Multiset [a] -> [[Multiset a]]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Multiset [a] -> [([a], Int)]
forall a. Multiset a -> [(a, Int)]
toCounts

uncollate :: ([a], Count) -> [(a, Count)]
uncollate :: forall a. ([a], Int) -> [(a, Int)]
uncollate ([a]
xs, Int
n) = (a -> (a, Int)) -> [a] -> [(a, Int)]
forall a b. (a -> b) -> [a] -> [b]
map (\a
x -> (a
x,Int
n)) [a]
xs