module Math.TreeFun.Tree where
import Data.List
import Data.Tree
import qualified Data.Map as M
import Data.Maybe
import qualified Data.Sequence as S
import Control.Applicative
import qualified Data.Foldable as F
import Control.Monad.State
import Math.TreeFun.Types
boolToInt :: Bool -> Int
boolToInt True = 1
boolToInt False = 0
isLeaf :: Tree a -> Bool
isLeaf (Node { subForest = [] }) = True
isLeaf _ = False
leaves :: Tree a -> [a]
leaves (Node { rootLabel = x, subForest = [] }) = [x]
leaves (Node { rootLabel = _, subForest = xs }) = concatMap leaves xs
leavesHeight :: (Ord a) => Int -> Tree a -> M.Map a Int
leavesHeight !h (Node { rootLabel = x, subForest = [] }) = M.singleton x h
leavesHeight !h (Node { rootLabel = _, subForest = xs }) =
M.unions . map (leavesHeight (h + 1)) $ xs
leavesCommonHeight :: (Ord a) => Int -> Tree a -> M.Map a (Int, Int)
leavesCommonHeight startHeight tree = evalState (iter startHeight tree) 0
where
iter !h (Node { rootLabel = x, subForest = [] }) = do
label <- get
return $ M.singleton x (h, label)
iter !h (Node { rootLabel = _, subForest = xs }) = do
ls <- mapM (iter (h + 1)) . filter isLeaf $ xs
label <- get
put $ label + 1
ts <- mapM (iter (h + 1)) . filter (not . isLeaf) $ xs
return . M.unions . (++) ts $ ls
leavesParentMult :: (Ord a) => Double
-> Double
-> Tree a
-> M.Map a (Double, Double)
leavesParentMult !w !d (Node { rootLabel = x, subForest = [] }) =
M.singleton x (w, d)
leavesParentMult !w !d (Node { rootLabel = _, subForest = xs }) =
M.unions . map (leavesParentMult (w * genericLength xs) (d + 1)) $ xs
leavesCommonParentMult :: (Ord a) => Int -> Tree a -> M.Map a (Int, Int)
leavesCommonParentMult numChildren tree = evalState (iter numChildren tree) 0
where
iter multChildren (Node { rootLabel = x, subForest = [] }) = do
label <- get
return $ M.singleton x (multChildren, label)
iter multChildren (Node { rootLabel = _, subForest = xs }) = do
ls <- mapM (iter (multChildren * length xs)) . filter isLeaf $ xs
label <- get
put $ label + 1
ts <- mapM (iter (multChildren * length xs))
. filter (not . isLeaf)
$ xs
return . M.unions . (++) ts $ ls
leavesHeightList :: Int -> Tree a -> [(a, Int)]
leavesHeightList h (Node { rootLabel = x, subForest = [] }) = [(x, h)]
leavesHeightList h (Node { rootLabel = _, subForest = xs }) =
concatMap (leavesHeightList (h + 1)) xs
innerNodes :: Tree a -> [a]
innerNodes (Node { rootLabel = _, subForest = [] }) = []
innerNodes (Node { rootLabel = x, subForest = xs }) = x
: concatMap innerNodes xs
numLeaves :: (Num b) => Tree a -> b
numLeaves = genericLength . leaves
numInner :: (Num b) => Tree a -> b
numInner = genericLength . innerNodes
hasRootLeaf :: Tree a -> Bool
hasRootLeaf (Node { subForest = ts }) = not . null . filter isLeaf $ ts
getRootLeaves :: Tree a -> [a]
getRootLeaves (Node { subForest = ts }) = map rootLabel . filter isLeaf $ ts
getProperties :: (Eq b) => PropertyMap a b -> [b]
getProperties = nub . F.toList . F.foldl' (S.><) S.empty . M.elems
filterLeaves :: Tree a -> Tree a
filterLeaves tree = tree {subForest = filter (not . isLeaf) . subForest $ tree}
filterRootLeaves :: Tree a -> Tree a
filterRootLeaves root@(Node { subForest = ts }) =
root { subForest = filter (not . isLeaf) ts }
getDistanceMap :: (Eq a, Ord a) => Tree a -> DistanceMap a
getDistanceMap tree = M.fromListWith (M.unionWith (S.><))
$ (\x y -> if x == y
then (x, M.singleton 0 (S.singleton y))
else ( x
, M.singleton
(getDistance tree x y)
(S.singleton y) ) )
<$> leaves tree
<*> leaves tree
getDistance :: (Eq a) => Tree a -> a -> a -> Int
getDistance (Node { rootLabel = l, subForest = [] }) x y = boolToInt
$ l `elem` [x, y]
getDistance n@(Node { rootLabel = _, subForest = xs }) x y
| none = 0
| otherwise = sum
. (:) (boolToInt notShared)
. map (\t -> getDistance t x y)
$ xs
where
notShared = (elem x ls) || (elem y ls) && not (elem x ls && elem y ls)
where
ls = leaves n
none = not (elem x ls || elem y ls)
where
ls = leaves n
getDistanceMapSuperNode :: (Eq a, Ord a) => Tree (SuperNode a) -> DistanceMap a
getDistanceMapSuperNode tree = M.fromListWith (M.unionWith (S.><))
$ (\x y -> if x == y
then
(x , M.singleton 0 (S.singleton y))
else ( x
, M.singleton
(getDistanceSuperNode tree x y)
(S.singleton y) ) )
<$> allLeaves
<*> allLeaves
where
allLeaves = M.keys . myLeaves . rootLabel $ tree
getDistanceSuperNode :: (Eq a, Ord a) => Tree (SuperNode a) -> a -> a -> Int
getDistanceSuperNode (Node { rootLabel = SuperNode { myLeaves = ls
, myParent = p }
, subForest = ts } ) x y
| shared ls = head
. filter (/= 1)
. map (\a -> getDistanceSuperNode a x y)
$ ts
| notShared ls = getParentLeafDist x p + getParentLeafDist y p
| otherwise = 0
where
notShared xs = (M.member x xs || M.member y xs)
&& not (M.member x xs && M.member y xs)
shared xs = M.member x xs && M.member y xs
getParentLeafDist a b = fst . fromJust . M.lookup a . myLeaves $ b
sumTree :: (Num a) => Tree a -> a
sumTree = F.foldl' (+) 0
toSuperNodeTree :: (Ord a) => SuperNode a -> Tree a -> Tree (SuperNode a)
toSuperNodeTree p n@(Node { rootLabel = x, subForest = xs }) =
Node { rootLabel = newNode
, subForest = map (toSuperNodeTree newNode) xs }
where
newNode = SuperNode { myRootLabel = x
, myLeaves = leavesCommonHeight 0 n
, myParent = p }