{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
module Math.Combinatorics.Multiset
(
Count
, Multiset(..)
, emptyMS, singletonMS
, consMS, (+:)
, toList
, fromList
, fromListEq
, fromDistinctList
, fromCounts
, getCounts
, size
, disjUnion
, disjUnions
, permutations
, permutationsRLE
, Vec
, vPartitions
, partitions
, splits
, kSubsets
, cycles
, bracelets
, genFixedBracelets
, 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
newtype Multiset a = MS { forall a. Multiset a -> [(a, Int)]
toCounts :: [(a, Count)] }
deriving (Int -> Multiset a -> ShowS
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 -> 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)
fromCounts :: [(a, Count)] -> Multiset a
fromCounts :: forall a. [(a, Int)] -> Multiset a
fromCounts = forall a. [(a, Int)] -> Multiset a
MS
getCounts :: Multiset a -> [Count]
getCounts :: forall a. Multiset a -> [Int]
getCounts = forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> b
snd forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Multiset a -> [(a, Int)]
toCounts
size :: Multiset a -> Int
size :: forall a. Multiset a -> Int
size = forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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) = forall a. [(a, Int)] -> Multiset a
MS ([(a, Int)] -> [(b, Int)]
f [(a, Int)]
m)
emptyMS :: Multiset a
emptyMS :: forall a. Multiset a
emptyMS = forall a. [(a, Int)] -> Multiset a
MS []
singletonMS :: a -> Multiset a
singletonMS :: forall a. a -> Multiset a
singletonMS a
a = forall a. [(a, Int)] -> Multiset a
MS [(a
a,Int
1)]
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 forall a. Ord a => a -> a -> Bool
> Int
0 = forall a. [(a, Int)] -> Multiset a
MS ((a, Int)
eforall a. a -> [a] -> [a]
:[(a, Int)]
m)
| Bool
otherwise = forall a. [(a, Int)] -> Multiset a
MS [(a, Int)]
m
(+:) :: (a, Count) -> Multiset a -> Multiset a
+: :: forall a. (a, Int) -> Multiset a -> Multiset a
(+:) = forall a. (a, Int) -> Multiset a -> Multiset a
consMS
toList :: Multiset a -> [a]
toList :: forall a. Multiset a -> [a]
toList = forall a. [(a, Int)] -> [a]
expandCounts forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Multiset a -> [(a, Int)]
toCounts
expandCounts :: [(a, Count)] -> [a]
expandCounts :: forall a. [(a, Int)] -> [a]
expandCounts = forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry (forall a b c. (a -> b -> c) -> b -> a -> c
flip forall a. Int -> a -> [a]
replicate))
fromList :: Ord a => [a] -> Multiset a
fromList :: forall a. Ord a => [a] -> Multiset a
fromList = forall a. [(a, Int)] -> Multiset a
fromCounts forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> b) -> [a] -> [b]
map (forall a. [a] -> a
head forall (a :: * -> * -> *) b c c'.
Arrow a =>
a b c -> a b c' -> a b (c, c')
&&& forall (t :: * -> *) a. Foldable t => t a -> Int
length) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Eq a => [a] -> [[a]]
group forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Ord a => [a] -> [a]
sort
fromListEq :: Eq a => [a] -> Multiset a
fromListEq :: forall a. Eq a => [a] -> Multiset a
fromListEq = forall a. [(a, Int)] -> Multiset a
fromCounts forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall {a}. Eq a => [a] -> [(a, Int)]
fromListEq'
where fromListEq' :: [a] -> [(a, Int)]
fromListEq' [] = []
fromListEq' (a
x:[a]
xs) = (a
x, Int
1 forall a. Num a => a -> a -> a
+ forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
xEqs) forall a. a -> [a] -> [a]
: [a] -> [(a, Int)]
fromListEq' [a]
xNeqs
where
([a]
xEqs, [a]
xNeqs) = forall a. (a -> Bool) -> [a] -> ([a], [a])
partition (forall a. Eq a => a -> a -> Bool
==a
x) [a]
xs
fromDistinctList :: [a] -> Multiset a
fromDistinctList :: forall a. [a] -> Multiset a
fromDistinctList = forall a. [(a, Int)] -> Multiset a
fromCounts forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> b) -> [a] -> [b]
map (\a
x -> (a
x,Int
1))
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) = forall a. [(a, Int)] -> Multiset a
MS ([(a, Int)]
xs forall a. [a] -> [a] -> [a]
++ [(a, Int)]
ys)
disjUnions :: [Multiset a] -> Multiset a
disjUnions :: forall a. [Multiset a] -> Multiset a
disjUnions = forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr forall a. Multiset a -> Multiset a -> Multiset a
disjUnion (forall a. [(a, Int)] -> Multiset a
MS [])
data RMultiset a = RMS (Maybe (a, Count)) [(a,Count)]
deriving Int -> RMultiset a -> ShowS
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
toRMS :: Multiset a -> RMultiset a
toRMS :: forall a. Multiset a -> RMultiset a
toRMS = forall a. Maybe (a, Int) -> [(a, Int)] -> RMultiset a
RMS forall a. Maybe a
Nothing forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Multiset a -> [(a, Int)]
toCounts
fromRMS :: RMultiset a -> Multiset a
fromRMS :: forall a. RMultiset a -> Multiset a
fromRMS (RMS Maybe (a, Int)
Nothing [(a, Int)]
m) = forall a. [(a, Int)] -> Multiset a
MS [(a, Int)]
m
fromRMS (RMS (Just (a, Int)
e) [(a, Int)]
m) = forall a. [(a, Int)] -> Multiset a
MS ((a, Int)
eforall a. a -> [a] -> [a]
:[(a, Int)]
m)
permutations :: Multiset a -> [[a]]
permutations :: forall a. Multiset a -> [[a]]
permutations = forall a b. (a -> b) -> [a] -> [b]
map forall a. [(a, Int)] -> [a]
expandCounts forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Multiset a -> [[(a, Int)]]
permutationsRLE
permutationsRLE :: Multiset a -> [[(a,Count)]]
permutationsRLE :: forall a. Multiset a -> [[(a, Int)]]
permutationsRLE (MS []) = [[]]
permutationsRLE Multiset a
m = forall a. RMultiset a -> [[(a, Int)]]
permutationsRLE' (forall a. Multiset a -> RMultiset a
toRMS Multiset a
m)
permutationsRLE' :: RMultiset a -> [[(a,Count)]]
permutationsRLE' :: forall a. RMultiset a -> [[(a, Int)]]
permutationsRLE' (RMS Maybe (a, Int)
Nothing [(a
x,Int
n)]) = [[(a
x,Int
n)]]
permutationsRLE' RMultiset a
m = [ (a, Int)
e forall a. a -> [a] -> [a]
: [(a, Int)]
p
| ((a, Int)
e, RMultiset a
m') <- forall a. RMultiset a -> [((a, Int), RMultiset a)]
selectRMS RMultiset a
m
, [(a, Int)]
p <- forall a. RMultiset a -> [[(a, Int)]]
permutationsRLE' RMultiset a
m'
]
selectRMS :: RMultiset a -> [((a, Count), RMultiset a)]
selectRMS :: forall a. RMultiset a -> [((a, Int), RMultiset a)]
selectRMS (RMS Maybe (a, Int)
_ []) = []
selectRMS (RMS Maybe (a, Int)
e ((a
x,Int
n) : [(a, Int)]
ms)) =
((a
x,Int
n), forall a. Maybe (a, Int) -> [(a, Int)] -> RMultiset a
RMS forall a. Maybe a
Nothing (forall b a. b -> (a -> b) -> Maybe a -> b
maybe [(a, Int)]
ms (forall a. a -> [a] -> [a]
:[(a, Int)]
ms) Maybe (a, Int)
e)) forall a. a -> [a] -> [a]
:
[ ( (a
x,Int
k), forall a. Maybe (a, Int) -> [(a, Int)] -> RMultiset a
RMS (forall a. a -> Maybe a
Just (a
x,Int
nforall a. Num a => a -> a -> a
-Int
k))
(forall b a. b -> (a -> b) -> Maybe a -> b
maybe [(a, Int)]
ms (forall a. a -> [a] -> [a]
:[(a, Int)]
ms) Maybe (a, Int)
e) )
| Int
k <- [Int
nforall a. Num a => a -> a -> a
-Int
1, Int
nforall a. Num a => a -> a -> a
-Int
2 .. Int
1]
] forall a. [a] -> [a] -> [a]
++
forall a b. (a -> b) -> [a] -> [b]
map (forall (a :: * -> * -> *) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second (forall a. (a, Int) -> RMultiset a -> RMultiset a
consRMS (a
x,Int
n))) (forall a. RMultiset a -> [((a, Int), RMultiset a)]
selectRMS (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) = forall a. Maybe (a, Int) -> [(a, Int)] -> RMultiset a
RMS Maybe (a, Int)
e ((a, Int)
xforall a. a -> [a] -> [a]
:[(a, Int)]
m)
type Vec = [Count]
(<|=) :: Vec -> Vec -> Bool
[Int]
xs <|= :: [Int] -> [Int] -> Bool
<|= [Int]
ys = forall (t :: * -> *). Foldable t => t Bool -> Bool
and forall a b. (a -> b) -> a -> b
$ forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith forall a. Ord a => a -> a -> Bool
(<=) [Int]
xs [Int]
ys
vZero :: Vec -> Vec
vZero :: [Int] -> [Int]
vZero = forall a b. (a -> b) -> [a] -> [b]
map (forall a b. a -> b -> a
const Int
0)
vIsZero :: Vec -> Bool
vIsZero :: [Int] -> Bool
vIsZero = forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (forall a. Eq a => a -> a -> Bool
==Int
0)
(.+.), (.-.) :: Vec -> Vec -> Vec
.+. :: [Int] -> [Int] -> [Int]
(.+.) = forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith forall a. Num a => a -> a -> a
(+)
.-. :: [Int] -> [Int] -> [Int]
(.-.) = forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (-)
(*.) :: Count -> Vec -> Vec
*. :: Int -> [Int] -> [Int]
(*.) Int
n = forall a b. (a -> b) -> [a] -> [b]
map (Int
nforall a. Num a => a -> a -> a
*)
vDiv :: Vec -> Vec -> Count
vDiv :: [Int] -> [Int] -> Int
vDiv [Int]
v1 [Int]
v2 = forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
minimum forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. [Maybe a] -> [a]
catMaybes forall a b. (a -> b) -> a -> b
$ forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith forall {a}. Integral a => a -> a -> Maybe a
zdiv [Int]
v1 [Int]
v2
where zdiv :: a -> a -> Maybe a
zdiv a
_ a
0 = forall a. Maybe a
Nothing
zdiv a
x a
y = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ a
x forall a. Integral a => a -> a -> a
`div` a
y
vInc :: Vec -> Vec -> Vec
vInc :: [Int] -> [Int] -> [Int]
vInc [Int]
lim [Int]
v = forall a. [a] -> [a]
reverse (forall {a}. (Num a, Ord a) => [a] -> [a] -> [a]
vInc' (forall a. [a] -> [a]
reverse [Int]
lim) (forall a. [a] -> [a]
reverse [Int]
v))
where vInc' :: [a] -> [a] -> [a]
vInc' [a]
_ [] = []
vInc' [] (a
x:[a]
xs) = a
xforall a. Num a => a -> a -> a
+a
1 forall a. a -> [a] -> [a]
: [a]
xs
vInc' (a
l:[a]
ls) (a
x:[a]
xs) | a
x forall a. Ord a => a -> a -> Bool
< a
l = a
xforall a. Num a => a -> a -> a
+a
1 forall a. a -> [a] -> [a]
: [a]
xs
| Bool
otherwise = a
0 forall a. a -> [a] -> [a]
: [a] -> [a] -> [a]
vInc' [a]
ls [a]
xs
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 = [forall a. [(a, Int)] -> Multiset a
MS []]
vPart [Int]
v [Int]
vL
| [Int]
v forall a. Ord a => a -> a -> Bool
<= [Int]
vL = []
| Bool
otherwise = forall a. [(a, Int)] -> Multiset a
MS [([Int]
v,Int
1)]
forall a. a -> [a] -> [a]
: [ ([Int]
v',Int
k) 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 :: Vec -> Vec
vHalf :: [Int] -> [Int]
vHalf [] = []
vHalf (Int
x:[Int]
xs) | (forall a. Integral a => a -> Bool
even Int
x) = (Int
x forall a. Integral a => a -> a -> a
`div` Int
2) forall a. a -> [a] -> [a]
: [Int] -> [Int]
vHalf [Int]
xs
| Bool
otherwise = (Int
x forall a. Integral a => a -> a -> a
`div` Int
2) forall a. a -> [a] -> [a]
: [Int]
xs
downFrom :: a -> [a]
downFrom a
n = [a
n,(a
nforall a. Num a => a -> a -> a
-a
1)..a
0]
within :: Vec -> [Vec]
within :: [Int] -> [[Int]]
within = forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
sequence forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> b) -> [a] -> [b]
map forall {a}. (Num a, Enum a) => a -> [a]
downFrom
clip :: Vec -> Vec -> Vec
clip :: [Int] -> [Int] -> [Int]
clip = forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith forall a. Ord a => a -> a -> a
min
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 forall a. Ord a => a -> a -> Bool
> [Int]
s = []
withinFromTo [Int]
m [Int]
s [Int]
e = 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
xforall a. a -> [a] -> [a]
:[a]
xs | a
x <- [a
start,(a
startforall a. Num a => a -> a -> a
-a
1)..a
end],
let useS' :: Bool
useS' = Bool
useS Bool -> Bool -> Bool
&& a
xforall a. Eq a => a -> a -> Bool
==a
s,
let useE' :: Bool
useE' = Bool
useE Bool -> Bool -> Bool
&& a
xforall 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' ]
partitions :: Multiset a -> [Multiset (Multiset a)]
partitions :: forall a. Multiset a -> [Multiset (Multiset a)]
partitions (MS []) = [forall a. [(a, Int)] -> Multiset a
MS []]
partitions (MS [(a, Int)]
m) = (forall a b. (a -> b) -> [a] -> [b]
map forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap) (forall {a}. [a] -> [Int] -> Multiset a
combine [a]
elts) forall a b. (a -> b) -> a -> b
$ [Int] -> [Multiset [Int]]
vPartitions [Int]
counts
where ([a]
elts, [Int]
counts) = forall a b. [(a, b)] -> ([a], [b])
unzip [(a, Int)]
m
combine :: [a] -> [Int] -> Multiset a
combine [a]
es [Int]
cs = forall a. [(a, Int)] -> Multiset a
MS forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. (a -> Bool) -> [a] -> [a]
filter ((forall a. Eq a => a -> a -> Bool
/=Int
0) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> b
snd) forall a b. (a -> b) -> a -> b
$ forall a b. [a] -> [b] -> [(a, b)]
zip [a]
es [Int]
cs
splits :: Multiset a -> [(Multiset a, Multiset a)]
splits :: forall a. Multiset a -> [(Multiset a, Multiset a)]
splits (MS []) = [(forall a. [(a, Int)] -> Multiset a
MS [], forall a. [(a, Int)] -> Multiset a
MS [])]
splits (MS ((a
x,Int
n):[(a, Int)]
m)) =
forall {a} {b}. [a] -> (a -> [b]) -> [b]
for [Int
0..Int
n] forall a b. (a -> b) -> a -> b
$ \Int
k ->
forall a b. (a -> b) -> [a] -> [b]
map (forall {a}. a -> Int -> Multiset a -> Multiset a
addElt a
x Int
k forall (a :: * -> * -> *) b c b' c'.
Arrow a =>
a b c -> a b' c' -> a (b, b') (c, c')
*** forall {a}. a -> Int -> Multiset a -> Multiset a
addElt a
x (Int
nforall a. Num a => a -> a -> a
-Int
k)) (forall a. Multiset a -> [(Multiset a, Multiset a)]
splits (forall a. [(a, Int)] -> Multiset a
MS [(a, Int)]
m))
kSubsets :: Count -> Multiset a -> [Multiset a]
kSubsets :: forall a. Int -> Multiset a -> [Multiset a]
kSubsets Int
0 Multiset a
_ = [forall a. [(a, Int)] -> Multiset a
MS []]
kSubsets Int
_ (MS []) = []
kSubsets Int
k (MS ((a
x,Int
n):[(a, Int)]
m)) =
forall {a} {b}. [a] -> (a -> [b]) -> [b]
for [Int
0 .. forall a. Ord a => a -> a -> a
min Int
k Int
n] forall a b. (a -> b) -> a -> b
$ \Int
j ->
forall a b. (a -> b) -> [a] -> [b]
map (forall {a}. a -> Int -> Multiset a -> Multiset a
addElt a
x Int
j) (forall a. Int -> Multiset a -> [Multiset a]
kSubsets (Int
k forall a. Num a => a -> a -> a
- Int
j) (forall a. [(a, Int)] -> Multiset a
MS [(a, Int)]
m))
for :: [a] -> (a -> [b]) -> [b]
for = forall a b c. (a -> b -> c) -> b -> a -> c
flip forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap
addElt :: a -> Int -> Multiset a -> Multiset a
addElt a
_ Int
0 = forall a. a -> a
id
addElt a
x Int
k = ((a
x,Int
k) forall a. (a, Int) -> Multiset a -> Multiset a
+:)
cycles :: Multiset a -> [[a]]
cycles :: forall a. Multiset a -> [[a]]
cycles (MS []) = []
cycles m :: Multiset a
m@(MS ((a
x1,Int
n1):[(a, Int)]
xs))
| Int
n1 forall a. Eq a => a -> a -> Bool
== Int
1 = (forall a.
Int -> Int -> Int -> [(Int, a)] -> [(Int, (a, Int))] -> [[a]]
cycles' Int
n Int
2 Int
1 [(Int
0,a
x1)] (forall a. [a] -> [a]
reverse forall a b. (a -> b) -> a -> b
$ forall a b. [a] -> [b] -> [(a, b)]
zip [Int
1..] [(a, Int)]
xs))
| Bool
otherwise = (forall a.
Int -> Int -> Int -> [(Int, a)] -> [(Int, (a, Int))] -> [[a]]
cycles' Int
n Int
2 Int
1 [(Int
0,a
x1)] (forall a. [a] -> [a]
reverse forall a b. (a -> b) -> a -> b
$ forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0..] ((a
x1,Int
n1forall a. Num a => a -> a -> a
-Int
1)forall a. a -> [a] -> [a]
:[(a, Int)]
xs)))
where n :: Int
n = forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Multiset a -> [Int]
getCounts forall a b. (a -> b) -> a -> b
$ Multiset a
m
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 forall a. Integral a => a -> a -> a
`mod` Int
p forall a. Eq a => a -> a -> Bool
== Int
0 = [forall a b. (a -> b) -> [a] -> [b]
map 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 =
(forall a. (a -> Bool) -> [a] -> [a]
takeWhile ((forall a. Ord a => a -> a -> Bool
>=Int
atp) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> a
fst) [(Int, (a, Int))]
xs) forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \(Int
j, (a
xj,Int
_)) ->
forall a.
Int -> Int -> Int -> [(Int, a)] -> [(Int, (a, Int))] -> [[a]]
cycles' Int
n (Int
tforall a. Num a => a -> a -> a
+Int
1) (if Int
j forall a. Eq a => a -> a -> Bool
== Int
atp then Int
p else Int
t)
((Int
j,a
xj)forall a. a -> [a] -> [a]
:[(Int, a)]
pre)
(forall a. Int -> [(Int, (a, Int))] -> [(Int, (a, Int))]
remove Int
j [(Int, (a, Int))]
xs)
where atp :: Int
atp = forall a b. (a, b) -> a
fst forall a b. (a -> b) -> a -> b
$ [(Int, a)]
pre forall a. [a] -> Int -> a
!! (Int
p 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 forall a. Eq a => a -> a -> Bool
== Int
j' Bool -> Bool -> Bool
&& Int
nj forall a. Eq a => a -> a -> Bool
== Int
1 = [(Int, (a, Int))]
xs
| Int
j forall a. Eq a => a -> a -> Bool
== Int
j' = (Int
j',(a
xj,Int
njforall a. Num a => a -> a -> a
-Int
1))forall a. a -> [a] -> [a]
:[(Int, (a, Int))]
xs
| Bool
otherwise = (Int, (a, Int))
xforall a. a -> [a] -> [a]
:forall a. Int -> [(Int, (a, Int))] -> [(Int, (a, Int))]
remove Int
j [(Int, (a, Int))]
xs
class Snocable p a where
(|>) :: p -> a -> p
class Indexable p where
(!) :: p -> Int -> Int
type PreNecklace = [Int]
data Pre = Pre !Int (Maybe Int) PreNecklace
deriving (Int -> Pre -> ShowS
[Pre] -> ShowS
Pre -> String
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 forall a. Maybe a
Nothing []
getPre :: Pre -> PreNecklace
getPre :: Pre -> [Int]
getPre (Pre Int
_ Maybe Int
_ [Int]
as) = 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 (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
tforall a. Num a => a -> a -> a
+Int
1) Maybe Int
a1 (Int
aforall 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 forall a. [a] -> Int -> a
!! (Int
tforall a. Num a => a -> a -> a
-Int
i)
data RLE a = RLE !Int !Int [(a,Int)]
deriving (Int -> RLE a -> ShowS
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)
emptyRLE :: RLE a
emptyRLE :: forall a. RLE a
emptyRLE = 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) forall a. Eq a => a -> a -> Bool
== (a
a2,Int
n2) = forall a. Ord a => [(a, Int)] -> [(a, Int)] -> Ordering
compareRLE [(a, Int)]
rle1 [(a, Int)]
rle2
| a
a1 forall a. Ord a => a -> a -> Bool
< a
a2 = Ordering
LT
| a
a1 forall a. Ord a => a -> a -> Bool
> a
a2 = Ordering
GT
| (Int
n1 forall a. Ord a => a -> a -> Bool
< Int
n2 Bool -> Bool -> Bool
&& (forall (t :: * -> *) a. Foldable t => t a -> Bool
null [(a, Int)]
rle1 Bool -> Bool -> Bool
|| forall a b. (a, b) -> a
fst (forall a. [a] -> a
head [(a, Int)]
rle1) forall a. Ord a => a -> a -> Bool
< a
a2)) Bool -> Bool -> Bool
|| (Int
n1 forall a. Ord a => a -> a -> Bool
> Int
n2 Bool -> Bool -> Bool
&& Bool -> Bool
not (forall (t :: * -> *) a. Foldable t => t a -> Bool
null [(a, Int)]
rle2) Bool -> Bool -> Bool
&& a
a1 forall a. Ord a => a -> a -> Bool
< forall a b. (a, b) -> a
fst (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
_ = 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 forall a. Ord a => a -> a -> Bool
<= Int
v = Int
a
| Bool
otherwise = (forall a. Int -> Int -> [(a, Int)] -> RLE a
RLE (Int
nforall a. Num a => a -> a -> a
-Int
v) (Int
bforall a. Num a => a -> a -> a
-Int
1) [(Int, Int)]
rest) forall p. Indexable p => p -> Int -> Int
! (Int
iforall 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' = 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 forall a. Eq a => a -> a -> Bool
== a
a' = forall a. Int -> Int -> [(a, Int)] -> RLE a
RLE (Int
nforall a. Num a => a -> a -> a
+Int
1) Int
b ((a
a,Int
vforall a. Num a => a -> a -> a
+Int
1)forall a. a -> [a] -> [a]
:[(a, Int)]
rest)
| Bool
otherwise = forall a. Int -> Int -> [(a, Int)] -> RLE a
RLE (Int
nforall a. Num a => a -> a -> a
+Int
1) (Int
bforall a. Num a => a -> a -> a
+Int
1) ((a
a',Int
1)forall a. a -> [a] -> [a]
:[(a, Int)]
rle)
data Pre' = Pre' Pre (RLE Int)
deriving Int -> Pre' -> ShowS
[Pre'] -> ShowS
Pre' -> String
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 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 forall p. Indexable p => p -> Int -> Int
! (Int
len forall a. Num a => a -> a -> a
- Int
i 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 forall p a. Snocable p a => p -> a -> p
|> Int
a) (RLE Int
rle forall p a. Snocable p a => p -> a -> p
|> Int
a)
type Bracelet = [Int]
genFixedBracelets :: Int -> [(Int,Int)] -> [Bracelet]
genFixedBracelets :: Int -> [(Int, Int)] -> [[Int]]
genFixedBracelets Int
n [(Int
0,Int
k)] | Int
k forall a. Ord a => a -> a -> Bool
>= Int
n = [forall a. Int -> a -> [a]
replicate Int
k Int
0]
| Bool
otherwise = []
genFixedBracelets Int
n [(Int, Int)]
content = forall w a. Writer w a -> w
execWriter (Int -> Int -> Int -> IntMap Int -> Pre' -> Writer [[Int]] ()
go Int
1 Int
1 Int
0 (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'
_ | forall a. IntMap a -> [Int]
IM.keys IntMap Int
con forall a. Eq a => a -> a -> Bool
== [Int
0] = 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 forall a. Ord a => a -> a -> Bool
> Int
n =
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (forall a. Int -> [a] -> [a]
take (Int
n forall a. Num a => a -> a -> a
- Int
r) [Int]
as forall a. Ord a => a -> a -> Bool
>= forall a. [a] -> [a]
reverse (forall a. Int -> [a] -> [a]
take (Int
nforall a. Num a => a -> a -> a
-Int
r) [Int]
as) Bool -> Bool -> Bool
&& Int
n forall a. Integral a => a -> a -> a
`mod` Int
p forall a. Eq a => a -> a -> Bool
== Int
0) forall a b. (a -> b) -> a -> b
$
forall (m :: * -> *) w. Monad m => w -> WriterT w m ()
tell [Pre' -> [Int]
getPre' Pre'
pre]
| Bool
otherwise = do
let a' :: Int
a' = Pre'
pre forall p. Indexable p => p -> Int -> Int
! (Int
tforall a. Num a => a -> a -> a
-Int
p)
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ (forall a. (a -> Bool) -> [a] -> [a]
dropWhile (forall a. Ord a => a -> a -> Bool
< Int
a') forall a b. (a -> b) -> a -> b
$ forall a. IntMap a -> [Int]
IM.keys IntMap Int
con) 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 forall p a. Snocable p a => p -> a -> p
|> Int
j
c :: Ordering
c = forall {p}. p -> Pre' -> Ordering
checkRev2 Int
t Pre'
pre'
p' :: Int
p' | Int
j forall a. Eq a => a -> a -> Bool
/= Int
a' = Int
t
| Bool
otherwise = Int
p
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Ordering
c forall a. Eq a => a -> a -> Bool
== Ordering
EQ) forall a b. (a -> b) -> a -> b
$ Int -> Int -> Int -> IntMap Int -> Pre' -> Writer [[Int]] ()
go (Int
tforall a. Num a => a -> a -> a
+Int
1) Int
p' Int
t IntMap Int
con' Pre'
pre'
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Ordering
c forall a. Eq a => a -> a -> Bool
== Ordering
GT) forall a b. (a -> b) -> a -> b
$ Int -> Int -> Int -> IntMap Int -> Pre' -> Writer [[Int]] ()
go (Int
tforall 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
| forall a. IntMap a -> Bool
IM.null IntMap Int
con = IntMap Int
con
| Bool
otherwise = forall a. (Maybe a -> Maybe a) -> Int -> IntMap a -> IntMap a
IM.alter 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) = forall a. Maybe a
Nothing
q (Just a
cnt) = forall a. a -> Maybe a
Just (a
cntforall a. Num a => a -> a -> a
-a
1)
q Maybe a
_ = forall a. Maybe a
Nothing
checkRev2 :: p -> Pre' -> Ordering
checkRev2 p
_ (Pre' Pre
_ (RLE Int
_ Int
_ [(Int, Int)]
rle)) = forall a. Ord a => [(a, Int)] -> [(a, Int)] -> Ordering
compareRLE [(Int, Int)]
rle (forall a. [a] -> [a]
reverse [(Int, Int)]
rle)
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 = forall a. [(Int, a)] -> IntMap a
IM.fromList (forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0..] (forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> a
fst [(a, Int)]
cnts))
content :: [(Int, Int)]
content = 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 (forall a. Multiset a -> Int
size Multiset a
ms) [(Int, Int)]
content
bs :: [[a]]
bs = forall a b. (a -> b) -> [a] -> [b]
map (forall a b. (a -> b) -> [a] -> [b]
map (forall a. HasCallStack => Maybe a -> a
fromJust forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b c. (a -> b -> c) -> b -> a -> c
flip forall a. Int -> IntMap a -> Maybe a
IM.lookup IntMap a
contentMap)) [[Int]]
rawBs
sequenceMS :: Multiset [a] -> [Multiset a]
sequenceMS :: forall a. Multiset [a] -> [Multiset a]
sequenceMS = forall a b. (a -> b) -> [a] -> [b]
map forall a. [Multiset a] -> Multiset a
disjUnions
forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
sequence
forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> b) -> [a] -> [b]
map (\([a]
xs, Int
n) -> forall a. Int -> Multiset a -> [Multiset a]
kSubsets Int
n (forall a. [(a, Int)] -> Multiset a
MS forall a b. (a -> b) -> a -> b
$ forall a. ([a], Int) -> [(a, Int)]
uncollate ([a]
xs, Int
n)))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Multiset a -> [(a, Int)]
toCounts
uncollate :: ([a], Count) -> [(a, Count)]
uncollate :: forall a. ([a], Int) -> [(a, Int)]
uncollate ([a]
xs, Int
n) = forall a b. (a -> b) -> [a] -> [b]
map (\a
x -> (a
x,Int
n)) [a]
xs