{-|
== PQ-tree

PQ-tree is a data structure for representing permutations. It's useful for solving the consecutive
ones problem.

This is an implementation of PQ-tree as described by Booth and Lueker. The tree consists of three
types of nodes, P, Q and leaf. Leaves are single elements. The frontier of the tree is the left to
right order of leaves in the tree. A P-node is equivalent to another with its children permuted in
any order. A Q-node is equivalent to one with its children reversed. The permutations represented by
a tree are the frontiers of all equivalent trees.

The various reduction templates described in Booth and Lueker's paper are used, but the construction
algorithm is not. This implementation is purely functional, hence it uses a simpler top-down
approach instead that runs in O(n) time (plus set member check overheads) rather than in
O(update size).
This implementation is specialized to Ints, but can be modified to work with other types.

Sources:

* https://en.wikipedia.org/wiki/PQ_tree
* Kellogg S. Booth and George S. Lueker, "Testing for the consecutive ones property, interval
  graphs, and graph planarity using PQ-tree algorithms", 1976
  https://www.sciencedirect.com/science/article/pii/S0022000076800451

-}

{-
Implementation notes:
* In reduce, the root of the pertinent subtree is found using pertinent node counts, then another
  pass is made through the pertinent subtree performing the actual reduction.
* The RNode type is a temporary node indicating empty, full, or partial status. In a reduction that
  does not reduce the tree to Nothing, there are at most two partial RNodes that float up to the
  root of the pertinent subtree.
-}

module PQTree
    ( PQNode
    , buildPQ
    , reducePQ
    , reduceAllPQ
    , frontierPQ
    , permsPQ
    ) where

import Control.Applicative
import Control.DeepSeq
import Control.Monad
import Control.Monad.State
import Data.Either
import Data.List
import Data.Maybe
import qualified Data.IntSet as IS

data PQNode = PQLeaf !Int
            | PNode  [PQNode]
            | QNode  [PQNode]
            deriving Int -> PQNode -> ShowS
[PQNode] -> ShowS
PQNode -> String
(Int -> PQNode -> ShowS)
-> (PQNode -> String) -> ([PQNode] -> ShowS) -> Show PQNode
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [PQNode] -> ShowS
$cshowList :: [PQNode] -> ShowS
show :: PQNode -> String
$cshow :: PQNode -> String
showsPrec :: Int -> PQNode -> ShowS
$cshowsPrec :: Int -> PQNode -> ShowS
Show

type Parts = ([PQNode], [PQNode]) -- empty and full, end first
data RNode = Empty PQNode
           | Full  PQNode
           | Part  Parts
           deriving Int -> RNode -> ShowS
[RNode] -> ShowS
RNode -> String
(Int -> RNode -> ShowS)
-> (RNode -> String) -> ([RNode] -> ShowS) -> Show RNode
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [RNode] -> ShowS
$cshowList :: [RNode] -> ShowS
show :: RNode -> String
$cshow :: RNode -> String
showsPrec :: Int -> RNode -> ShowS
$cshowsPrec :: Int -> RNode -> ShowS
Show

mkPNode, mkQNode :: [PQNode] -> PQNode
mkPNode :: [PQNode] -> PQNode
mkPNode []     = String -> PQNode
forall a. HasCallStack => String -> a
error String
"empty PNode"
mkPNode [PQNode
n]    = PQNode
n
mkPNode [PQNode]
ns     = [PQNode] -> PQNode
PNode [PQNode]
ns
mkQNode :: [PQNode] -> PQNode
mkQNode []     = String -> PQNode
forall a. HasCallStack => String -> a
error String
"empty QNode"
mkQNode [PQNode
n]    = PQNode
n
mkQNode [PQNode
n, PQNode
m] = [PQNode] -> PQNode
PNode [PQNode
n, PQNode
m]
mkQNode [PQNode]
ns     = [PQNode] -> PQNode
QNode [PQNode]
ns

mkPartial :: [PQNode] -> [PQNode] -> RNode
mkPartial :: [PQNode] -> [PQNode] -> RNode
mkPartial [] [PQNode]
_  = String -> RNode
forall a. HasCallStack => String -> a
error String
"empty RNode es"
mkPartial [PQNode]
_  [] = String -> RNode
forall a. HasCallStack => String -> a
error String
"empty RNode fs"
mkPartial [PQNode]
es [PQNode]
fs = Parts -> RNode
Part ([PQNode]
es, [PQNode]
fs)

emptyMany, fullMany :: State [RNode] [PQNode]
emptyMany :: State [RNode] [PQNode]
emptyMany = ([RNode] -> ([PQNode], [RNode])) -> State [RNode] [PQNode]
forall s (m :: * -> *) a. MonadState s m => (s -> (a, s)) -> m a
state (([RNode] -> ([PQNode], [RNode])) -> State [RNode] [PQNode])
-> ([RNode] -> ([PQNode], [RNode])) -> State [RNode] [PQNode]
forall a b. (a -> b) -> a -> b
$ [PQNode] -> [RNode] -> ([PQNode], [RNode])
go [] where
    go :: [PQNode] -> [RNode] -> ([PQNode], [RNode])
go [PQNode]
ys (Empty PQNode
n : [RNode]
xs) = [PQNode] -> [RNode] -> ([PQNode], [RNode])
go (PQNode
nPQNode -> [PQNode] -> [PQNode]
forall a. a -> [a] -> [a]
:[PQNode]
ys) [RNode]
xs
    go [PQNode]
ys [RNode]
xs             = ([PQNode] -> [PQNode]
forall a. [a] -> [a]
reverse [PQNode]
ys, [RNode]
xs)
fullMany :: State [RNode] [PQNode]
fullMany = ([RNode] -> ([PQNode], [RNode])) -> State [RNode] [PQNode]
forall s (m :: * -> *) a. MonadState s m => (s -> (a, s)) -> m a
state (([RNode] -> ([PQNode], [RNode])) -> State [RNode] [PQNode])
-> ([RNode] -> ([PQNode], [RNode])) -> State [RNode] [PQNode]
forall a b. (a -> b) -> a -> b
$ [PQNode] -> [RNode] -> ([PQNode], [RNode])
go [] where
    go :: [PQNode] -> [RNode] -> ([PQNode], [RNode])
go [PQNode]
ys (Full PQNode
n : [RNode]
xs) = [PQNode] -> [RNode] -> ([PQNode], [RNode])
go (PQNode
nPQNode -> [PQNode] -> [PQNode]
forall a. a -> [a] -> [a]
:[PQNode]
ys) [RNode]
xs
    go [PQNode]
ys [RNode]
xs            = ([PQNode] -> [PQNode]
forall a. [a] -> [a]
reverse [PQNode]
ys, [RNode]
xs)

partialMaybe :: State [RNode] (Maybe Parts)
partialMaybe :: State [RNode] (Maybe Parts)
partialMaybe = ([RNode] -> (Maybe Parts, [RNode])) -> State [RNode] (Maybe Parts)
forall s (m :: * -> *) a. MonadState s m => (s -> (a, s)) -> m a
state [RNode] -> (Maybe Parts, [RNode])
go where
    go :: [RNode] -> (Maybe Parts, [RNode])
go (Part Parts
p : [RNode]
xs) = (Parts -> Maybe Parts
forall a. a -> Maybe a
Just Parts
p,  [RNode]
xs)
    go [RNode]
xs            = (Maybe Parts
forall a. Maybe a
Nothing, [RNode]
xs)

splitForP :: [RNode] -> ([PQNode], [Parts], [PQNode])
splitForP :: [RNode] -> ([PQNode], [Parts], [PQNode])
splitForP [RNode]
us = [RNode]
-> [PQNode] -> [Parts] -> [PQNode] -> ([PQNode], [Parts], [PQNode])
go [RNode]
us [] [] [] where
    go :: [RNode]
-> [PQNode] -> [Parts] -> [PQNode] -> ([PQNode], [Parts], [PQNode])
go []     [PQNode]
es [Parts]
ps [PQNode]
fs = ([PQNode]
es, [Parts]
ps, [PQNode]
fs)
    go (RNode
x:[RNode]
xs) [PQNode]
es [Parts]
ps [PQNode]
fs = case RNode
x of
        Empty PQNode
n -> [RNode]
-> [PQNode] -> [Parts] -> [PQNode] -> ([PQNode], [Parts], [PQNode])
go [RNode]
xs (PQNode
nPQNode -> [PQNode] -> [PQNode]
forall a. a -> [a] -> [a]
:[PQNode]
es) [Parts]
ps     [PQNode]
fs
        Part Parts
p  -> [RNode]
-> [PQNode] -> [Parts] -> [PQNode] -> ([PQNode], [Parts], [PQNode])
go [RNode]
xs [PQNode]
es     (Parts
pParts -> [Parts] -> [Parts]
forall a. a -> [a] -> [a]
:[Parts]
ps) [PQNode]
fs
        Full PQNode
n  -> [RNode]
-> [PQNode] -> [Parts] -> [PQNode] -> ([PQNode], [Parts], [PQNode])
go [RNode]
xs [PQNode]
es     [Parts]
ps     (PQNode
nPQNode -> [PQNode] -> [PQNode]
forall a. a -> [a] -> [a]
:[PQNode]
fs)

splitForQ :: [RNode] -> Maybe ([PQNode], Maybe Parts, [PQNode])
splitForQ :: [RNode] -> Maybe ([PQNode], Maybe Parts, [PQNode])
splitForQ [RNode]
xs = [RNode] -> Maybe ([PQNode], Maybe Parts, [PQNode])
go [RNode]
xs Maybe ([PQNode], Maybe Parts, [PQNode])
-> Maybe ([PQNode], Maybe Parts, [PQNode])
-> Maybe ([PQNode], Maybe Parts, [PQNode])
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> [RNode] -> Maybe ([PQNode], Maybe Parts, [PQNode])
go ([RNode] -> [RNode]
forall a. [a] -> [a]
reverse [RNode]
xs) where
    go :: [RNode] -> Maybe ([PQNode], Maybe Parts, [PQNode])
go = State [RNode] (Maybe ([PQNode], Maybe Parts, [PQNode]))
-> [RNode] -> Maybe ([PQNode], Maybe Parts, [PQNode])
forall s a. State s a -> s -> a
evalState (State [RNode] (Maybe ([PQNode], Maybe Parts, [PQNode]))
 -> [RNode] -> Maybe ([PQNode], Maybe Parts, [PQNode]))
-> State [RNode] (Maybe ([PQNode], Maybe Parts, [PQNode]))
-> [RNode]
-> Maybe ([PQNode], Maybe Parts, [PQNode])
forall a b. (a -> b) -> a -> b
$ [PQNode]
-> Maybe Parts
-> [PQNode]
-> Bool
-> Maybe ([PQNode], Maybe Parts, [PQNode])
forall (f :: * -> *) a b c.
Alternative f =>
a -> b -> c -> Bool -> f (a, b, c)
go' ([PQNode]
 -> Maybe Parts
 -> [PQNode]
 -> Bool
 -> Maybe ([PQNode], Maybe Parts, [PQNode]))
-> State [RNode] [PQNode]
-> StateT
     [RNode]
     Identity
     (Maybe Parts
      -> [PQNode] -> Bool -> Maybe ([PQNode], Maybe Parts, [PQNode]))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> State [RNode] [PQNode]
emptyMany StateT
  [RNode]
  Identity
  (Maybe Parts
   -> [PQNode] -> Bool -> Maybe ([PQNode], Maybe Parts, [PQNode]))
-> State [RNode] (Maybe Parts)
-> StateT
     [RNode]
     Identity
     ([PQNode] -> Bool -> Maybe ([PQNode], Maybe Parts, [PQNode]))
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> State [RNode] (Maybe Parts)
partialMaybe StateT
  [RNode]
  Identity
  ([PQNode] -> Bool -> Maybe ([PQNode], Maybe Parts, [PQNode]))
-> State [RNode] [PQNode]
-> StateT
     [RNode] Identity (Bool -> Maybe ([PQNode], Maybe Parts, [PQNode]))
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> State [RNode] [PQNode]
fullMany StateT
  [RNode] Identity (Bool -> Maybe ([PQNode], Maybe Parts, [PQNode]))
-> StateT [RNode] Identity Bool
-> State [RNode] (Maybe ([PQNode], Maybe Parts, [PQNode]))
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> ([RNode] -> Bool) -> StateT [RNode] Identity Bool
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets [RNode] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null
    go' :: a -> b -> c -> Bool -> f (a, b, c)
go' a
es b
ps c
fs Bool
end = (a
es, b
ps, c
fs) (a, b, c) -> f () -> f (a, b, c)
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ Bool -> f ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard Bool
end

splitForQRoot :: [RNode] -> Maybe ([PQNode], Maybe Parts, [PQNode], Maybe Parts, [PQNode])
splitForQRoot :: [RNode]
-> Maybe ([PQNode], Maybe Parts, [PQNode], Maybe Parts, [PQNode])
splitForQRoot = State
  [RNode]
  (Maybe ([PQNode], Maybe Parts, [PQNode], Maybe Parts, [PQNode]))
-> [RNode]
-> Maybe ([PQNode], Maybe Parts, [PQNode], Maybe Parts, [PQNode])
forall s a. State s a -> s -> a
evalState (State
   [RNode]
   (Maybe ([PQNode], Maybe Parts, [PQNode], Maybe Parts, [PQNode]))
 -> [RNode]
 -> Maybe ([PQNode], Maybe Parts, [PQNode], Maybe Parts, [PQNode]))
-> State
     [RNode]
     (Maybe ([PQNode], Maybe Parts, [PQNode], Maybe Parts, [PQNode]))
-> [RNode]
-> Maybe ([PQNode], Maybe Parts, [PQNode], Maybe Parts, [PQNode])
forall a b. (a -> b) -> a -> b
$
    [PQNode]
-> Maybe Parts
-> [PQNode]
-> Maybe Parts
-> [PQNode]
-> Bool
-> Maybe ([PQNode], Maybe Parts, [PQNode], Maybe Parts, [PQNode])
forall (f :: * -> *) a b c d e.
Alternative f =>
a -> b -> c -> d -> e -> Bool -> f (a, b, c, d, e)
go ([PQNode]
 -> Maybe Parts
 -> [PQNode]
 -> Maybe Parts
 -> [PQNode]
 -> Bool
 -> Maybe ([PQNode], Maybe Parts, [PQNode], Maybe Parts, [PQNode]))
-> State [RNode] [PQNode]
-> StateT
     [RNode]
     Identity
     (Maybe Parts
      -> [PQNode]
      -> Maybe Parts
      -> [PQNode]
      -> Bool
      -> Maybe ([PQNode], Maybe Parts, [PQNode], Maybe Parts, [PQNode]))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> State [RNode] [PQNode]
emptyMany StateT
  [RNode]
  Identity
  (Maybe Parts
   -> [PQNode]
   -> Maybe Parts
   -> [PQNode]
   -> Bool
   -> Maybe ([PQNode], Maybe Parts, [PQNode], Maybe Parts, [PQNode]))
-> State [RNode] (Maybe Parts)
-> StateT
     [RNode]
     Identity
     ([PQNode]
      -> Maybe Parts
      -> [PQNode]
      -> Bool
      -> Maybe ([PQNode], Maybe Parts, [PQNode], Maybe Parts, [PQNode]))
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> State [RNode] (Maybe Parts)
partialMaybe StateT
  [RNode]
  Identity
  ([PQNode]
   -> Maybe Parts
   -> [PQNode]
   -> Bool
   -> Maybe ([PQNode], Maybe Parts, [PQNode], Maybe Parts, [PQNode]))
-> State [RNode] [PQNode]
-> StateT
     [RNode]
     Identity
     (Maybe Parts
      -> [PQNode]
      -> Bool
      -> Maybe ([PQNode], Maybe Parts, [PQNode], Maybe Parts, [PQNode]))
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> State [RNode] [PQNode]
fullMany StateT
  [RNode]
  Identity
  (Maybe Parts
   -> [PQNode]
   -> Bool
   -> Maybe ([PQNode], Maybe Parts, [PQNode], Maybe Parts, [PQNode]))
-> State [RNode] (Maybe Parts)
-> StateT
     [RNode]
     Identity
     ([PQNode]
      -> Bool
      -> Maybe ([PQNode], Maybe Parts, [PQNode], Maybe Parts, [PQNode]))
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> State [RNode] (Maybe Parts)
partialMaybe StateT
  [RNode]
  Identity
  ([PQNode]
   -> Bool
   -> Maybe ([PQNode], Maybe Parts, [PQNode], Maybe Parts, [PQNode]))
-> State [RNode] [PQNode]
-> StateT
     [RNode]
     Identity
     (Bool
      -> Maybe ([PQNode], Maybe Parts, [PQNode], Maybe Parts, [PQNode]))
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> State [RNode] [PQNode]
emptyMany StateT
  [RNode]
  Identity
  (Bool
   -> Maybe ([PQNode], Maybe Parts, [PQNode], Maybe Parts, [PQNode]))
-> StateT [RNode] Identity Bool
-> State
     [RNode]
     (Maybe ([PQNode], Maybe Parts, [PQNode], Maybe Parts, [PQNode]))
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> ([RNode] -> Bool) -> StateT [RNode] Identity Bool
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets [RNode] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null
  where
    go :: a -> b -> c -> d -> e -> Bool -> f (a, b, c, d, e)
go a
es1 b
ps1 c
fs d
ps2 e
es2 Bool
end = (a
es1, b
ps1, c
fs, d
ps2, e
es2) (a, b, c, d, e) -> f () -> f (a, b, c, d, e)
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ Bool -> f ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard Bool
end

-- | Builds a PQ-tree from a set of Ints. O(n).
buildPQ :: [Int] -> PQNode
buildPQ :: [Int] -> PQNode
buildPQ = [PQNode] -> PQNode
mkPNode ([PQNode] -> PQNode) -> ([Int] -> [PQNode]) -> [Int] -> PQNode
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int -> PQNode) -> [Int] -> [PQNode]
forall a b. (a -> b) -> [a] -> [b]
map Int -> PQNode
PQLeaf

-- | Reduces a PQ-tree with a set of Ints. This should be a subset of what it was constructed with. O(n).
reducePQ :: [Int] -> PQNode -> Maybe PQNode
reducePQ :: [Int] -> PQNode -> Maybe PQNode
reducePQ []  = PQNode -> Maybe PQNode
forall a. a -> Maybe a
Just
reducePQ [Int]
xs0 = Maybe PQNode -> Either Int (Maybe PQNode) -> Maybe PQNode
forall b a. b -> Either a b -> b
fromRight (String -> Maybe PQNode
forall a. HasCallStack => String -> a
error String
"outside initial set") (Either Int (Maybe PQNode) -> Maybe PQNode)
-> (PQNode -> Either Int (Maybe PQNode)) -> PQNode -> Maybe PQNode
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PQNode -> Either Int (Maybe PQNode)
visit where
    xs :: IntSet
xs = [Int] -> IntSet
IS.fromList [Int]
xs0
    xsz :: Int
xsz = IntSet -> Int
IS.size IntSet
xs

    visit :: PQNode -> Either Int (Maybe PQNode)
    visit :: PQNode -> Either Int (Maybe PQNode)
visit n :: PQNode
n@(PNode [PQNode]
cs) = PQNode
-> ([PQNode] -> PQNode) -> [PQNode] -> Either Int (Maybe PQNode)
visitPQ PQNode
n [PQNode] -> PQNode
PNode [PQNode]
cs
    visit n :: PQNode
n@(QNode [PQNode]
cs) = PQNode
-> ([PQNode] -> PQNode) -> [PQNode] -> Either Int (Maybe PQNode)
visitPQ PQNode
n [PQNode] -> PQNode
QNode [PQNode]
cs
    visit n :: PQNode
n@(PQLeaf Int
x)
        | Int
x Int -> IntSet -> Bool
`IS.notMember` IntSet
xs = Int -> Either Int (Maybe PQNode)
forall a b. a -> Either a b
Left Int
0
        | Int
xsz Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
1             = Int -> Either Int (Maybe PQNode)
forall a b. a -> Either a b
Left Int
1
        | Bool
otherwise           = Maybe PQNode -> Either Int (Maybe PQNode)
forall a b. b -> Either a b
Right (Maybe PQNode -> Either Int (Maybe PQNode))
-> Maybe PQNode -> Either Int (Maybe PQNode)
forall a b. (a -> b) -> a -> b
$ PQNode -> Maybe PQNode
reduceRoot PQNode
n
    visitPQ :: PQNode
-> ([PQNode] -> PQNode) -> [PQNode] -> Either Int (Maybe PQNode)
visitPQ PQNode
n [PQNode] -> PQNode
f [PQNode]
cs
        | Int
cnt Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
xsz  = Maybe PQNode -> Either Int (Maybe PQNode)
forall a b. b -> Either a b
Right (Maybe PQNode -> Either Int (Maybe PQNode))
-> Maybe PQNode -> Either Int (Maybe PQNode)
forall a b. (a -> b) -> a -> b
$ PQNode -> Maybe PQNode
reduceRoot PQNode
n
        | [Maybe PQNode] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Maybe PQNode]
rts    = Int -> Either Int (Maybe PQNode)
forall a b. a -> Either a b
Left Int
cnt
        | ~[Maybe PQNode
r] <- [Maybe PQNode]
rts = Maybe PQNode -> Either Int (Maybe PQNode)
forall a b. b -> Either a b
Right (Maybe PQNode -> Either Int (Maybe PQNode))
-> Maybe PQNode -> Either Int (Maybe PQNode)
forall a b. (a -> b) -> a -> b
$ [PQNode] -> PQNode
f ([PQNode] -> PQNode) -> (PQNode -> [PQNode]) -> PQNode -> PQNode
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PQNode -> [PQNode]
getcs' (PQNode -> PQNode) -> Maybe PQNode -> Maybe PQNode
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe PQNode
r
      where
        ys :: [Either Int (Maybe PQNode)]
ys = (PQNode -> Either Int (Maybe PQNode))
-> [PQNode] -> [Either Int (Maybe PQNode)]
forall a b. (a -> b) -> [a] -> [b]
map PQNode -> Either Int (Maybe PQNode)
visit [PQNode]
cs
        cnt :: Int
cnt = [Int] -> Int
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ([Int] -> Int) -> [Int] -> Int
forall a b. (a -> b) -> a -> b
$ [Either Int (Maybe PQNode)] -> [Int]
forall a b. [Either a b] -> [a]
lefts [Either Int (Maybe PQNode)]
ys
        rts :: [Maybe PQNode]
rts = [Either Int (Maybe PQNode)] -> [Maybe PQNode]
forall a b. [Either a b] -> [b]
rights [Either Int (Maybe PQNode)]
ys
        getcs' :: PQNode -> [PQNode]
getcs' PQNode
c' = (PQNode -> Either Int (Maybe PQNode) -> PQNode)
-> [PQNode] -> [Either Int (Maybe PQNode)] -> [PQNode]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\PQNode
c -> (Int -> PQNode)
-> (Maybe PQNode -> PQNode) -> Either Int (Maybe PQNode) -> PQNode
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either (PQNode -> Int -> PQNode
forall a b. a -> b -> a
const PQNode
c) (PQNode -> Maybe PQNode -> PQNode
forall a b. a -> b -> a
const PQNode
c')) [PQNode]
cs [Either Int (Maybe PQNode)]
ys

    reduceRoot :: PQNode -> Maybe PQNode
    reduceRoot :: PQNode -> Maybe PQNode
reduceRoot n :: PQNode
n@(PNode [PQNode]
cs) = ([PQNode], [Parts], [PQNode]) -> Maybe PQNode
go (([PQNode], [Parts], [PQNode]) -> Maybe PQNode)
-> ([RNode] -> ([PQNode], [Parts], [PQNode]))
-> [RNode]
-> Maybe PQNode
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [RNode] -> ([PQNode], [Parts], [PQNode])
splitForP ([RNode] -> Maybe PQNode) -> Maybe [RNode] -> Maybe PQNode
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< (PQNode -> Maybe RNode) -> [PQNode] -> Maybe [RNode]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM PQNode -> Maybe RNode
reduceInternal [PQNode]
cs where
        go :: ([PQNode], [Parts], [PQNode]) -> Maybe PQNode
go ([PQNode]
es, [Parts]
ps, [PQNode]
fs) = case [Parts]
ps of
            []       -> PQNode -> Maybe PQNode
forall a. a -> Maybe a
Just PQNode
noPartial
            [Parts
p1]     -> PQNode -> Maybe PQNode
forall a. a -> Maybe a
Just (PQNode -> Maybe PQNode) -> PQNode -> Maybe PQNode
forall a b. (a -> b) -> a -> b
$ Parts -> Parts -> PQNode
withPartial Parts
p1 ([], [])
            [Parts
p1, Parts
p2] -> PQNode -> Maybe PQNode
forall a. a -> Maybe a
Just (PQNode -> Maybe PQNode) -> PQNode -> Maybe PQNode
forall a b. (a -> b) -> a -> b
$ Parts -> Parts -> PQNode
withPartial Parts
p1 Parts
p2
            [Parts]
_        -> Maybe PQNode
forall a. Maybe a
Nothing
          where
            noPartial :: PQNode
noPartial
                | [PQNode] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [PQNode]
es Bool -> Bool -> Bool
|| [PQNode] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [PQNode]
fs = PQNode
n
                | Bool
otherwise          = [PQNode] -> PQNode
mkPNode ([PQNode] -> PQNode) -> [PQNode] -> PQNode
forall a b. (a -> b) -> a -> b
$ [PQNode] -> PQNode
mkPNode [PQNode]
fs PQNode -> [PQNode] -> [PQNode]
forall a. a -> [a] -> [a]
: [PQNode]
es
            withPartial :: Parts -> Parts -> PQNode
withPartial ([PQNode]
pe1, [PQNode]
pf1) ([PQNode]
pe2, [PQNode]
pf2) = [PQNode] -> PQNode
mkPNode ([PQNode] -> PQNode) -> [PQNode] -> PQNode
forall a b. (a -> b) -> a -> b
$ PQNode
qn PQNode -> [PQNode] -> [PQNode]
forall a. a -> [a] -> [a]
: [PQNode]
es where
                qn :: PQNode
qn = [PQNode] -> PQNode
mkQNode ([PQNode] -> PQNode) -> [PQNode] -> PQNode
forall a b. (a -> b) -> a -> b
$
                    [PQNode]
pe1 [PQNode] -> [PQNode] -> [PQNode]
forall a. [a] -> [a] -> [a]
++ [PQNode] -> [PQNode]
forall a. [a] -> [a]
reverse [PQNode]
pf1 [PQNode] -> [PQNode] -> [PQNode]
forall a. [a] -> [a] -> [a]
++ [[PQNode] -> PQNode
mkPNode [PQNode]
fs | Bool -> Bool
not ([PQNode] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [PQNode]
fs)] [PQNode] -> [PQNode] -> [PQNode]
forall a. [a] -> [a] -> [a]
++ [PQNode]
pf2 [PQNode] -> [PQNode] -> [PQNode]
forall a. [a] -> [a] -> [a]
++ [PQNode] -> [PQNode]
forall a. [a] -> [a]
reverse [PQNode]
pe2

    reduceRoot n :: PQNode
n@(QNode [PQNode]
cs) = (([PQNode], Maybe Parts, [PQNode], Maybe Parts, [PQNode])
 -> PQNode)
-> Maybe ([PQNode], Maybe Parts, [PQNode], Maybe Parts, [PQNode])
-> Maybe PQNode
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ([PQNode], Maybe Parts, [PQNode], Maybe Parts, [PQNode]) -> PQNode
go (Maybe ([PQNode], Maybe Parts, [PQNode], Maybe Parts, [PQNode])
 -> Maybe PQNode)
-> ([RNode]
    -> Maybe ([PQNode], Maybe Parts, [PQNode], Maybe Parts, [PQNode]))
-> [RNode]
-> Maybe PQNode
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [RNode]
-> Maybe ([PQNode], Maybe Parts, [PQNode], Maybe Parts, [PQNode])
splitForQRoot ([RNode] -> Maybe PQNode) -> Maybe [RNode] -> Maybe PQNode
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< (PQNode -> Maybe RNode) -> [PQNode] -> Maybe [RNode]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM PQNode -> Maybe RNode
reduceInternal [PQNode]
cs where
        go :: ([PQNode], Maybe Parts, [PQNode], Maybe Parts, [PQNode]) -> PQNode
go ([PQNode]
es1, Maybe Parts
ps1, [PQNode]
fs, Maybe Parts
ps2, [PQNode]
es2) = case (Maybe Parts
ps1, Maybe Parts
ps2) of
            (Maybe Parts
Nothing, Maybe Parts
Nothing) -> PQNode
n
            (Maybe Parts, Maybe Parts)
_                  -> Parts -> Parts -> PQNode
withPartial (Parts -> Maybe Parts -> Parts
forall a. a -> Maybe a -> a
fromMaybe ([], []) Maybe Parts
ps1) (Parts -> Maybe Parts -> Parts
forall a. a -> Maybe a -> a
fromMaybe ([], []) Maybe Parts
ps2)
          where
            withPartial :: Parts -> Parts -> PQNode
withPartial ([PQNode]
pe1, [PQNode]
pf1) ([PQNode]
pe2, [PQNode]
pf2) =
                [PQNode] -> PQNode
mkQNode ([PQNode] -> PQNode) -> [PQNode] -> PQNode
forall a b. (a -> b) -> a -> b
$ [PQNode]
es1 [PQNode] -> [PQNode] -> [PQNode]
forall a. [a] -> [a] -> [a]
++ [PQNode]
pe1 [PQNode] -> [PQNode] -> [PQNode]
forall a. [a] -> [a] -> [a]
++ [PQNode] -> [PQNode]
forall a. [a] -> [a]
reverse [PQNode]
pf1 [PQNode] -> [PQNode] -> [PQNode]
forall a. [a] -> [a] -> [a]
++ [PQNode]
fs [PQNode] -> [PQNode] -> [PQNode]
forall a. [a] -> [a] -> [a]
++ [PQNode]
pf2 [PQNode] -> [PQNode] -> [PQNode]
forall a. [a] -> [a] -> [a]
++ [PQNode] -> [PQNode]
forall a. [a] -> [a]
reverse [PQNode]
pe2 [PQNode] -> [PQNode] -> [PQNode]
forall a. [a] -> [a] -> [a]
++ [PQNode]
es2

    reduceRoot n :: PQNode
n@(PQLeaf Int
_) = PQNode -> Maybe PQNode
forall a. a -> Maybe a
Just PQNode
n

    reduceInternal :: PQNode -> Maybe RNode
    reduceInternal :: PQNode -> Maybe RNode
reduceInternal n :: PQNode
n@(PNode [PQNode]
cs) = ([PQNode], [Parts], [PQNode]) -> Maybe RNode
go (([PQNode], [Parts], [PQNode]) -> Maybe RNode)
-> ([RNode] -> ([PQNode], [Parts], [PQNode]))
-> [RNode]
-> Maybe RNode
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [RNode] -> ([PQNode], [Parts], [PQNode])
splitForP ([RNode] -> Maybe RNode) -> Maybe [RNode] -> Maybe RNode
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< (PQNode -> Maybe RNode) -> [PQNode] -> Maybe [RNode]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM PQNode -> Maybe RNode
reduceInternal [PQNode]
cs where
        go :: ([PQNode], [Parts], [PQNode]) -> Maybe RNode
go ([PQNode]
es, [Parts]
ps, [PQNode]
fs) = case [Parts]
ps of
            []   -> RNode -> Maybe RNode
forall a. a -> Maybe a
Just RNode
noPartial
            [Parts
p1] -> RNode -> Maybe RNode
forall a. a -> Maybe a
Just (RNode -> Maybe RNode) -> RNode -> Maybe RNode
forall a b. (a -> b) -> a -> b
$ Parts -> RNode
withPartial Parts
p1
            [Parts]
_    -> Maybe RNode
forall a. Maybe a
Nothing
          where
            noPartial :: RNode
noPartial
                | [PQNode] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [PQNode]
es   = PQNode -> RNode
Full PQNode
n
                | [PQNode] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [PQNode]
fs   = PQNode -> RNode
Empty PQNode
n
                | Bool
otherwise = [PQNode] -> [PQNode] -> RNode
mkPartial [[PQNode] -> PQNode
mkPNode [PQNode]
es] [[PQNode] -> PQNode
mkPNode [PQNode]
fs]
            withPartial :: Parts -> RNode
withPartial ([PQNode]
pe, [PQNode]
pf) = [PQNode] -> [PQNode] -> RNode
mkPartial
                ([[PQNode] -> PQNode
mkPNode [PQNode]
es | Bool -> Bool
not ([PQNode] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [PQNode]
es)] [PQNode] -> [PQNode] -> [PQNode]
forall a. [a] -> [a] -> [a]
++ [PQNode]
pe) ([[PQNode] -> PQNode
mkPNode [PQNode]
fs | Bool -> Bool
not ([PQNode] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [PQNode]
fs)] [PQNode] -> [PQNode] -> [PQNode]
forall a. [a] -> [a] -> [a]
++ [PQNode]
pf)

    reduceInternal n :: PQNode
n@(QNode [PQNode]
cs) = (([PQNode], Maybe Parts, [PQNode]) -> RNode)
-> Maybe ([PQNode], Maybe Parts, [PQNode]) -> Maybe RNode
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ([PQNode], Maybe Parts, [PQNode]) -> RNode
go (Maybe ([PQNode], Maybe Parts, [PQNode]) -> Maybe RNode)
-> ([RNode] -> Maybe ([PQNode], Maybe Parts, [PQNode]))
-> [RNode]
-> Maybe RNode
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [RNode] -> Maybe ([PQNode], Maybe Parts, [PQNode])
splitForQ ([RNode] -> Maybe RNode) -> Maybe [RNode] -> Maybe RNode
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< (PQNode -> Maybe RNode) -> [PQNode] -> Maybe [RNode]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM PQNode -> Maybe RNode
reduceInternal [PQNode]
cs where
        go :: ([PQNode], Maybe Parts, [PQNode]) -> RNode
go ([PQNode]
es, Maybe Parts
ps, [PQNode]
fs) = RNode -> (Parts -> RNode) -> Maybe Parts -> RNode
forall b a. b -> (a -> b) -> Maybe a -> b
maybe RNode
noPartial Parts -> RNode
withPartial Maybe Parts
ps where
            noPartial :: RNode
noPartial
                | [PQNode] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [PQNode]
es   = PQNode -> RNode
Full PQNode
n
                | [PQNode] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [PQNode]
fs   = PQNode -> RNode
Empty PQNode
n
                | Bool
otherwise = [PQNode] -> [PQNode] -> RNode
mkPartial [PQNode]
es ([PQNode] -> [PQNode]
forall a. [a] -> [a]
reverse [PQNode]
fs)
            withPartial :: Parts -> RNode
withPartial ([PQNode]
pe, [PQNode]
pf) = [PQNode] -> [PQNode] -> RNode
mkPartial ([PQNode]
es [PQNode] -> [PQNode] -> [PQNode]
forall a. [a] -> [a] -> [a]
++ [PQNode]
pe) ([PQNode] -> [PQNode]
forall a. [a] -> [a]
reverse [PQNode]
fs [PQNode] -> [PQNode] -> [PQNode]
forall a. [a] -> [a] -> [a]
++ [PQNode]
pf)

    reduceInternal n :: PQNode
n@(PQLeaf Int
x) = RNode -> Maybe RNode
forall a. a -> Maybe a
Just (RNode -> Maybe RNode) -> RNode -> Maybe RNode
forall a b. (a -> b) -> a -> b
$ if Int
x Int -> IntSet -> Bool
`IS.member` IntSet
xs then PQNode -> RNode
Full PQNode
n else PQNode -> RNode
Empty PQNode
n

-- | Reduces a PQ-tree with many sets. O(mn) for m sets.
reduceAllPQ :: [[Int]] -> PQNode -> Maybe PQNode
reduceAllPQ :: [[Int]] -> PQNode -> Maybe PQNode
reduceAllPQ [[Int]]
xss PQNode
t = (PQNode -> [Int] -> Maybe PQNode)
-> PQNode -> [[Int]] -> Maybe PQNode
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM (([Int] -> PQNode -> Maybe PQNode)
-> PQNode -> [Int] -> Maybe PQNode
forall a b c. (a -> b -> c) -> b -> a -> c
flip [Int] -> PQNode -> Maybe PQNode
reducePQ) PQNode
t [[Int]]
xss

-- | The frontier of the PQ-tree. O(n).
frontierPQ :: PQNode -> [Int]
frontierPQ :: PQNode -> [Int]
frontierPQ = (PQNode -> [Int] -> [Int]) -> [Int] -> PQNode -> [Int]
forall a b c. (a -> b -> c) -> b -> a -> c
flip PQNode -> [Int] -> [Int]
go [] where
    go :: PQNode -> [Int] -> [Int]
go (PQLeaf Int
x) [Int]
acc = Int
xInt -> [Int] -> [Int]
forall a. a -> [a] -> [a]
:[Int]
acc
    go (PNode [PQNode]
cs) [Int]
acc = (PQNode -> [Int] -> [Int]) -> [Int] -> [PQNode] -> [Int]
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr PQNode -> [Int] -> [Int]
go [Int]
acc [PQNode]
cs
    go (QNode [PQNode]
cs) [Int]
acc = (PQNode -> [Int] -> [Int]) -> [Int] -> [PQNode] -> [Int]
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr PQNode -> [Int] -> [Int]
go [Int]
acc [PQNode]
cs

-- | The permutations represented by the PQ-tree. O(do not use).
permsPQ :: PQNode -> [[Int]]
permsPQ :: PQNode -> [[Int]]
permsPQ PQNode
n = case PQNode
n of
    PQLeaf Int
x -> [[Int
x]]
    PNode [PQNode]
cs -> ([[[Int]]] -> [[Int]]) -> [[[[Int]]]] -> [[Int]]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap [[[Int]]] -> [[Int]]
forall a. [[[a]]] -> [[a]]
prod ([[[[Int]]]] -> [[Int]]) -> [[[[Int]]]] -> [[Int]]
forall a b. (a -> b) -> a -> b
$ [[[Int]]] -> [[[[Int]]]]
forall a. [a] -> [[a]]
permutations ([[[Int]]] -> [[[[Int]]]]) -> [[[Int]]] -> [[[[Int]]]]
forall a b. (a -> b) -> a -> b
$ (PQNode -> [[Int]]) -> [PQNode] -> [[[Int]]]
forall a b. (a -> b) -> [a] -> [b]
map PQNode -> [[Int]]
permsPQ [PQNode]
cs
    QNode [PQNode]
cs -> [[[Int]]] -> [[Int]]
forall a. [[[a]]] -> [[a]]
prod ((PQNode -> [[Int]]) -> [PQNode] -> [[[Int]]]
forall a b. (a -> b) -> [a] -> [b]
map PQNode -> [[Int]]
permsPQ [PQNode]
cs) [[Int]] -> [[Int]] -> [[Int]]
forall a. [a] -> [a] -> [a]
++ [[[Int]]] -> [[Int]]
forall a. [[[a]]] -> [[a]]
prod ([[[Int]]] -> [[[Int]]]
forall a. [a] -> [a]
reverse ([[[Int]]] -> [[[Int]]]) -> [[[Int]]] -> [[[Int]]]
forall a b. (a -> b) -> a -> b
$ (PQNode -> [[Int]]) -> [PQNode] -> [[[Int]]]
forall a b. (a -> b) -> [a] -> [b]
map PQNode -> [[Int]]
permsPQ [PQNode]
cs)
  where
    prod :: [[[a]]] -> [[a]]
prod = ([[a]] -> [a]) -> [[[a]]] -> [[a]]
forall a b. (a -> b) -> [a] -> [b]
map [[a]] -> [a]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([[[a]]] -> [[a]]) -> ([[[a]]] -> [[[a]]]) -> [[[a]]] -> [[a]]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [[[a]]] -> [[[a]]]
forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
sequence

-- Note: permsPQ is far from optimal, but it is unlikely to be used with larger values of n anyway

--------------------------------------------------------------------------------
-- For tests

instance NFData PQNode where
    rnf :: PQNode -> ()
rnf (PQLeaf Int
x) = Int -> ()
forall a. NFData a => a -> ()
rnf Int
x
    rnf (PNode [PQNode]
cs) = [PQNode] -> ()
forall a. NFData a => a -> ()
rnf [PQNode]
cs
    rnf (QNode [PQNode]
cs) = [PQNode] -> ()
forall a. NFData a => a -> ()
rnf [PQNode]
cs