Tree.hs

module Tree where

data Tree = Leaf Int | Node Int Tree Tree
  deriving (Show, Read, Eq, Ord)

type Forest = [Tree]

mkTree :: Int -> Tree
mkTree n =
  case n of
    0 -> Leaf n
    _ -> Node n (mkTree $ n-1) (mkTree $ n-1)

mkForest :: (Tree -> Bool) -> Bool -> Tree -> Forest -> Forest
mkForest f parentAdded tr acc =
  case tr of
    Leaf _ -> case (f tr, parentAdded) of
                (_, True) -> acc
                (False, False) -> acc ++ [tr]
                _ -> acc

    Node _ l r -> case (f tr, parentAdded) of
                    (True,_)       -> mkForest f False r (mkForest f False l acc)
                    (False, True)  -> mkForest f True r (mkForest f True l acc)
                    (False, False) -> let acc'  = acc ++ [tr]
                                      in mkForest f True r (mkForest f True l acc')

toOrderedTree :: Tree -> Tree
toOrderedTree = go 1
  where
    go n tr =
      case tr of
        Leaf _ -> Leaf n
        Node _ l r -> Node n (go (2 * n) l) (go (2 * n + 1) r)

printForest :: Forest -> [Int]
printForest fr = map (\tr -> case tr of
                           Leaf n -> n
                           Node n _ _ -> n)
                 fr

-- Make an erasure for `f`
mkErasure :: (Int -> Bool) -> Tree -> Bool
mkErasure f tr =
  case tr of
    Leaf n     -> f n
    Node n _ _ -> f n

tree1 :: Tree
tree1 = toOrderedTree $ mkTree 2

-- Node 1 (Node 2 (Leaf 4) (Leaf 5)) (Node 3 (Leaf 6) (Leaf 7))

test1 :: Bool
test1 = printForest (mkForest (mkErasure even) False tree1 []) == [1,5]

test2 :: Bool
test2 = printForest (mkForest (mkErasure odd) False tree1 []) == [2,6]

test3 :: Bool
test3 = printForest (mkForest (mkErasure (\n -> n `rem` 3 == 0)) False tree1 []) == [1,7]