Number all occurring leaves in a tree from left to right in Haskell

Question

The function type would be Tree a -> Tree (a, Int). I want to carry the count throughout the tree and number each occurring leaf accordingly.

So far I have tried this:

labelTree :: Tree a -> Tree (a, Int)
labelTree (Leaf a) = Leaf (a,1)
labelTree (tr)     = labelTree' (tr) 0

labelTree' :: Tree a -> Int -> (Tree (a,Int))
labelTree' (Leaf a) n   = Leaf (a,(n+1))
labelTree' (Node l r) n = (labelTree' (r) (snd (labelTree' (l) n)))

The problem is I am not sure why it is giving me the type error for this expression: labelTree' (Node l r) n = (labelTree' (r) (snd (labelTree' (l) n)))

Please indicate where I went wrong!

Answer 1

I have the same idea in mind as chepner: to use State. However, you don't have to write the recursion yourself because this is a simple traversal of the tree! Instead, derive Traversable and Foldable for your tree (good ideas anyway), and then lean on them to do the recursion for you:

{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE DeriveTraversable #-}

import qualified Control.Monad.Trans.State.Strict as S
data Tree a = Leaf a | Node (Tree a) (Tree a)
            deriving (Show, Functor, Foldable, Traversable)

labelTree :: Tree a -> Tree (a, Int)
labelTree t = S.evalState (traverse applyLabel t) 0
  where applyLabel x = do
          n <- S.get
          S.modify' succ
          pure (x, n)

*Main> labelTree (Node (Node (Leaf 'a') (Leaf 'b')) (Leaf 'c'))
Node (Node (Leaf ('a',0)) (Leaf ('b',1))) (Leaf ('c',2))

One nice feature of this implementation is that it will still work if you change the structure of your tree (e.g., to store data in the interior nodes). It is impossible to make mistakes like swapping the order of nodes, because you don't work at that level at all: Traversable handles it for you.