Haskell, list of lists from a tree

Question

I have this data structure for a tree:

data Tree a = NodeT a (Tree a) ( Tree a) | EmptyT

I need to create a function that returns a list of lists where each element of the list represents a level of the tree. For instance, from this:

          1
         / \
       2     3
      / \   / \
     4   5 6   7     

to this: [[1],[2,3],[4,5,6,7]]

The function must have the following form:

                     f :: Tree a -> [[a]]

How to do it using recursion?

anyone?

Thanks

Answer 1

Answer

levels :: Tree a -> [[a]]
levels t = levels' t []

levels' :: Tree a -> [[a]] -> [[a]]
levels' EmptyT rest = rest
levels' (NodeT a l r) [] = [a] : levels' l (levels r)
levels' (NodeT a l r) (x : xs) = (a : x) : levels' l (levels' r xs)

A slightly more complicated, but lazier, implementation of levels':

levels' EmptyT rest = rest
levels' (NodeT a l r) rest = (a : front) : levels' l (levels' r back)
  where
    (front, back) = case rest of
       [] -> ([], [])
       (x : xs) -> (x, xs)

Fans of folds will note that these are structured as catamorphisms:

cata :: (a -> b -> b -> b) -> b -> Tree a -> b
cata n e = go
  where
    go EmptyT = e
    go (NodeT a l r) = n a (go l) (go r)

levels t = cata br id t []
  where
    br a l r rest = (a : front) : l (r back)
      where
        (front, back) = case rest of
          [] -> ([], [])
          (x : xs) -> (x, xs)

As chi points out, there seems to be some connection between this general approach and the result of using Jakub Daniel's solution with difference lists as intermediate forms. This could look something like

import Data.Monoid

levels :: Tree a -> [[a]]
levels = map (flip appEndo []) . (cata br [])
  where
    br :: a -> [Endo [a]] -> [Endo [a]] -> [Endo [a]]
    br a l r = Endo (a :) : merge l r

merge :: Monoid a => [a] -> [a] -> [a]
merge [] ys = ys
merge (x : xs) ys = (x <> y) : merge xs ys'
   where
     (y,ys') =
       case ys of
         [] -> (mempty, [])
         p : ps -> (p, ps)

I'm not entirely sure just how this compares with the more direct approaches.

Discussion

Kostiantyn Rybnikov's answer cites Okasaki's Breadth-First Numbering: Lessons from a Small Exercise in Algorithm Design, an excellent paper that highlights many functional programmers' "blind spots" and offers good arguments for making abstract data types easy enough to use that they won't be missed. However, the problem that paper describes is significantly more complex than this one; not so much machinery is required here. Also, the paper notes that level-oriented solutions are actually slightly faster than queue-based ones in ML; I'd expect to see a larger difference in a lazy language like Haskell.

Jakub Daniel's answer attempts a level-oriented solution, but unfortunately has an efficiency problem. It builds each level by repeatedly appending one list to another, and those lists may all be of equal length. Thus in the worst case, if I am calculating this correctly, it takes O(n log n) to process a tree with n elements.

The approach I chose is level-oriented, but avoids the pain of concatenation by passing each left subtree the levels of its right sibling and cousins. Each node/leaf of the tree is processed exactly once. That processing involves O(1) work: pattern matching on that node/leaf, and, if it is a node, pattern matching on the list derived from the right sibling and cousins. Thus the total time is O(n) to process a tree with n elements.

Answer 2

You recursively compute the levels and always merge lists from two subtrees point-wise (thus all the slices in the same depth get merged together).

f :: Tree a -> [[a]]
f EmptyT = []
f (NodeT a t1 t2) = [a] : merge (f t1) (f t2)

merge :: [[a]] -> [[a]] -> [[a]]
merge [] ys = ys
merge xs [] = xs
merge (x:xs) (y:ys) = (x ++ y) : merge xs ys

If the tree were complete (all the paths from the root to a list are of the same length) then you could use zipWith (++) as merge.

Answer 3

Slightly more complicated solution, than the one which was accepted, but I think mine might be better in terms of memory consumption (it's a bit late, so please check yourself).

Intuition goes from a wonderful paper of Chris Okasaki "Breadth-First Numbering: Lessons from a Small Exercise in Algorithm Design". You can get general intuition on breadth-first tree traversals of trees in functional languages in detail.

I did somewhat ugly addition to add the "list of lists" splitting, there might be a better way:

module Main where

data Tree a = NodeT a (Tree a) (Tree a) | EmptyT

--      1
--     / \
--   2     3
--  / \   / \
-- 4   5 6   7     

f :: Tree a -> [[a]]
f t = joinBack (f' [(t, True)])

type UpLevel = Bool

f' :: [(Tree a, UpLevel)] -> [(a, UpLevel)]
f' [] = []
f' ((EmptyT, _) : ts) = f' ts
f' ((NodeT a t1 t2, up) : ts) = (a, up) : f' (ts ++ [(t1, up)] ++ [(t2, False)])

joinBack :: [(a, UpLevel)] -> [[a]]
joinBack = go []
  where
    go acc [] = [reverse acc]
    go acc ((x, False) : xs) = go (x : acc) xs
    go acc ((x, True) : xs) = reverse acc : go [] ((x, False):xs)

main :: IO ()
main = do
  let tree = NodeT 1 (NodeT 2 (NodeT 4 EmptyT EmptyT) (NodeT 5 EmptyT EmptyT))
                     (NodeT 3 (NodeT 6 EmptyT EmptyT) (NodeT 7 EmptyT EmptyT))
             :: Tree Int
  print (tail (f tree))