1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44
|
module Tree where
data Tree = Tree (Array Int [Int]) (Array Int (Int,Int,Int)) Int Int
-- Polymorphism here really should be handled with a class that has the members below
-- If we allow adding branches to functions later, we could move polymorphic definitions into files. e.g. for show.
data RootedTree = RootedTree Tree Int (Array Int Bool)
root (RootedTree _ r _) = r
rooted_tree t r = RootedTree t r (mkArray n check_away_from_root)
where check_away_from_root b = sourceNode t b == root t || or $ map (away_from_root t) (edgesBeforeEdge t)
n = numBranches t * 2
away_from_root (RootedTree t r arr) b = arr!b
-- For numNodes, numBranches, edgesOutOfNode, and nodesForEdge I'm currently using fake polymorphism
numNodes (Tree _ _ n _) = n
numNodes (RootedTree t _ _) = numNodes t
numBranches (Tree _ _ _ n) = n
numBranches (RootedTree t _ _) = numBranches t
edgesOutOfNode (Tree nodesArray _ _ _) node = nodesArray ! node
edgesOutOfNode (RootedTree t _ _) node = edgesOutOfNode t node
nodesForEdge (Tree _ branchesArray _ _) edgeIndex = branchesArray ! edgeIndex
nodesForEdge (RootedTree t _ _) edgeIndex = nodesForEdge t edgeIndex
sourceNode t b = let (s,_,_,_) = nodesForEdge t b in s
sourceIndex t b = let (_,i,_,_) = nodesForEdge t b in i
targetNode t b = let (_,_,t,_) = nodesForEdge t b in t
reverseEdge t b = let (_,_,_,r) = nodesForEdge t b in r
edgeForNodes t (n1,n2) = head [b | b <- (edgesOutOfNode t n1), (targetNode t b)==n2]
nodeDegree t n = length (edgesOutOfNode t n)
neighbors t n = fmap (targetNode t) (edgesOutOfNode t n)
edgesBeforeEdge t b = let (source,index,_,_) = nodesForEdge t b
in map (reverseEdge t) $ remove_element index $ edgesOutOfNode t source
is_leaf_node t n = (nodeDegree t n == 1)
is_internal_node t n = not $ is_leaf_node t n
nodes t = [0..numNodes t - 1]
leaf_nodes t = filter (is_leaf_node t) (nodes t)
internal_nodes t = filter (is_internal_node t) (nodes t)
remove_element _ [] = []
remove_element 0 (x:xs) = xs
remove_element i (x:xs) = x:(remove_element (i-1) xs)
|