WIP:

week_4/folding_with_trees.hs

@@ -8,10 +8,18 @@ foldTree :: [a] -> Tree a
 foldTree = foldr insertNode Leaf
 
 insertNode :: a -> Tree a -> Tree a
-insertNode n Leaf = Node 0 Leaf n Leaf
-insertNode n (Node _ Leaf x Leaf) = Node 1 (insertNode n Leaf) x Leaf
-insertNode n (Node _ lTree x Leaf) = Node 1 lTree x (insertNode n Leaf)
-insertNode n (Node _ Leaf x rTree) = Node 1 (insertNode n Leaf) x rTree
-insertNode n (Node height lTree@(Node lHeight _ _ _) x rTree@(Node rHeight _ _ _))
-  | lHeight > rHeight = Node (rHeight + 1) lTree x (insertNode n rTree)
-  | otherwise         = Node (lHeight + 1) (insertNode n lTree) x rTree
+insertNode x Leaf = Node 0 Leaf x Leaf
+insertNode x (Node _ Leaf v r) = Node 1 (insertNode x Leaf) v r
+insertNode x (Node _ l v Leaf) = Node 1 l v (insertNode x Leaf)
+insertNode x (Node h lTree@(Node lHeight _ _ _) y rTree@(Node rHeight _ _ _))
+  | lHeight > rHeight = insertLeft
+  | otherwise         = insertRight
+    where lHeight = heightOf lTree
+          rHeight = heightOf rTree
+          insertLeft  = Node (h + lHeight) (insertNode x lTree) y rTree
+          insertRight = Node (h + rHeight) lTree y (insertNode x rTree)
+
+heightOf :: Tree a -> Int
+heightOf Leaf           = undefined
+heightOf (Node h _ _ _) = h
+