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.