module ELynx.Data.Tree.Bipartition
(
Bipartition ()
, bp
, bpmap
, bipartitions
, bipartitionToBranch
, bipartitionsCombined
) where
import qualified Data.Map as M
import Data.Maybe
import qualified Data.Set as S
import Data.Tree
import ELynx.Data.Tree.Tree
newtype Bipartition a = Bipartition (S.Set a, S.Set a)
deriving (Show, Read)
bp :: Ord a => S.Set a -> S.Set a -> Bipartition a
bp x y = if x >= y
then Bipartition (x, y)
else Bipartition (y, x)
bpWith :: (Ord a, Ord b) => (a -> b) -> S.Set a -> S.Set a -> Bipartition b
bpWith f x y = bpmap f $ bp x y
instance (Eq a) => Eq (Bipartition a) where
Bipartition x == Bipartition y = x == y
instance (Ord a) => Ord (Bipartition a) where
Bipartition x `compare` Bipartition y = x `compare` y
bpmap :: (Ord a, Ord b) => (a -> b) -> Bipartition a -> Bipartition b
bpmap f (Bipartition (x, y)) = bp (S.map f x) (S.map f y)
leavesTree :: (Ord a) => Tree a -> Tree (S.Set a)
leavesTree (Node l []) = Node (S.singleton l) []
leavesTree (Node _ xs) = Node (S.unions $ map rootLabel xs') xs'
where xs' = map leavesTree xs
subForestGetLeafSets :: (Ord a)
=> S.Set a
-> Tree (S.Set a)
-> [S.Set a]
subForestGetLeafSets lvsS t = lvsOthers
where
xs = subForest t
nChildren = length xs
lvsChildren = map rootLabel xs
lvsOtherChildren = [ S.unions $ lvsS
: take i lvsChildren ++ drop (i+1) lvsChildren
| i <- [0 .. (nChildren - 1)] ]
lvsOthers = map (S.union lvsS) lvsOtherChildren
bipartitions :: Ord a => Tree a -> S.Set (Bipartition a)
bipartitions t = if S.size (S.fromList lvs) == length lvs
then bipartitionsUnsafe t
else error "bipartitions: The tree contains duplicate leaves."
where lvs = leaves t
bipartitionsUnsafe :: Ord a => Tree a -> S.Set (Bipartition a)
bipartitionsUnsafe (Node _ [] ) = S.empty
bipartitionsUnsafe (Node _ [x]) = bipartitionsUnsafe x
bipartitionsUnsafe t =
S.unions [ bipartitions' lvs x
| (lvs, x) <- zip lvsOthers (subForest lvsTree) ]
where
lvsTree = leavesTree t
lvsOthers = subForestGetLeafSets S.empty lvsTree
bipartitions' :: Ord a => S.Set a -> Tree (S.Set a) -> S.Set (Bipartition a)
bipartitions' lvsStem t@(Node lvs xs)
| S.null lvsStem = error "bipartitions': no complementing leaf set."
| null xs = S.singleton $ bp lvsStem lvs
| length xs == 1 = bipartitions' lvsStem (head xs)
| otherwise = S.unions $ S.singleton (bp lvsStem lvs) : zipWith bipartitions' lvsOthers xs
where
lvsOthers = subForestGetLeafSets lvsStem t
bipartitionToBranch :: (Ord a, Ord b, Monoid c)
=> (a -> b)
-> (a -> c)
-> Tree a
-> M.Map (Bipartition b) c
bipartitionToBranch f g t = if S.size (S.fromList lvs) == length lvs
then bipartitionToBranchUnsafe f g t
else error "bipartitionToBranch: The tree contains duplicate leaves."
where lvs = leaves t
bipartitionToBranchUnsafe :: (Ord a, Ord b, Monoid c)
=> (a -> b)
-> (a -> c)
-> Tree a
-> M.Map (Bipartition b) c
bipartitionToBranchUnsafe _ _ (Node _ [] ) = M.empty
bipartitionToBranchUnsafe f g (Node _ [x]) = bipartitionToBranchUnsafe f g x
bipartitionToBranchUnsafe f g t =
M.unionsWith (<>) [ bipartitionToBranch' lvs mempty f g x
| (lvs, x) <- zip lvsOthers (subForest nodeAndLeavesTrees) ]
where
lvsTree = leavesTree t
nodeAndLeavesTrees = fromJust $ merge t lvsTree
lvsOthers = subForestGetLeafSets S.empty lvsTree
bipartitionToBranch' :: (Ord a, Ord b, Monoid c)
=> S.Set a
-> c
-> (a -> b)
-> (a -> c)
-> Tree (a, S.Set a)
-> M.Map (Bipartition b) c
bipartitionToBranch' lvsStem br f g t@(Node l xs )
| S.null lvsStem = error "bipartitionToBranch': no complementing leaf set."
| null xs = M.singleton (bpWith f lvsStem lvsThisNode) (br <> g label)
| length xs == 1 = bipartitionToBranch' lvsStem (br <> g label) f g (head xs)
| otherwise = M.insert (bpWith f lvsStem lvsThisNode) (br <> g label)
$ M.unions [ bipartitionToBranch' lvs mempty f g x
| (lvs, x) <- zip lvsOthers xs ]
where
label = fst l
lvsThisNode = snd l
lvsOthers = subForestGetLeafSets lvsStem $ fmap snd t
bipartitionsCombined :: (Ord a, Show a) => Tree a -> S.Set (Bipartition a)
bipartitionsCombined t@(Node _ xs)
| null xs = S.empty
| length xs == 1 = bipartitionsCombined (head xs)
| length xs > 3 = S.empty
| otherwise = res
where
res = S.unions [ bipartitionsCombined' lvs x
| (lvs, x) <- zip lvsOthers (subForest lvsTree) ]
lvsTree = leavesTree t
lvsOthers = subForestGetLeafSets S.empty lvsTree
bipartitionsCombined' :: Ord a => S.Set a -> Tree (S.Set a) -> S.Set (Bipartition a)
bipartitionsCombined' lvsStem t@(Node lvs xs)
| S.null lvsStem = error "bipartitionsCombined': no complementing leaf set."
| null xs = S.singleton $ bp lvsStem lvs
| length xs == 1 = bipartitionsCombined' lvsStem (head xs)
| length xs == 2 = S.unions $
S.singleton (bp lvsStem lvs) : zipWith bipartitionsCombined' lvsOthers xs
| otherwise = S.singleton $ bp lvsStem lvs
where
lvsOthers = subForestGetLeafSets lvsStem t