Tuesday, July 5, 2011

Build a binary tree from pre-order traversal and in-order traversal

This post comes from i has 1337 code by way of An Algorithm a Day. I'm really just posting this to show the power of Scala collections.

The idea is that you have two sequential representations of a tree: one for in-order traversal, and the order for pre-order traversal. Just one of these representations are not enough to reconstruct the tree, but the two of them, in the absence of duplicate elements, is. See the links above for an example.

The problem lends itself to recursive solutions, but there's some list manipulation required to build the input for the recursive steps. The key point of the algorithm is the realization that the first element of pre-order is the root, and every element to the left of said root in the in-order belongs to the left of the tree, and every element to the right belongs to the right of the tree. The rest is pretty trivial.

Even so, Scala makes the trivial, well, trivial. Here's the full code:

case class Tree[T](el: T, left: Option[Tree[T]], right: Option[Tree[T]])

def mkTree[T](preorder: List[T], inorder: List[T]): Option[Tree[T]] = preorder.headOption map { head =>
    val (left, _ :: right) = inorder span (head !=)
    Tree(head, 
         mkTree(preorder filter (left contains), left), 
         mkTree(preorder filter (right contains), right))
}

Note: I'm using List instead of Seq so that I can use the :: extractor.

5 comments:

  1. Elegant!
    The realization of your understanding of the problem is sublime Scala

    ReplyDelete
  2. From Richard Carlsson erlang-questions%40erlang.org. a lazy list:
    LazyList :: () -> [] | [term() | LazyList]

    ReplyDelete
  3. I don't know scala, but what is the behavior of your program if the preorder(and/or in-order)list is ill-formed ?

    ReplyDelete
  4. That depends on what you mean by ill-formed. Most variations simply represent a valid tree. If, however, there are elements in one list nor present in the other, there can be two results. Elements from in order missing from preorder will be ignored. Elements from preorder missing from in order will result in an exception being thrown.

    ReplyDelete
  5. well done sir, well done

    ReplyDelete