Iterative Binary Tree Traversal

Traversal of binary trees has a simple recursive solution. This blog describes an iterative C++ function derived from a defunctionalised Ocaml solution to the problem.

Introduction
Consider the following code to compute the height of a binary tree.

struct Tree{
  Tree* left;
  Tree* right;
  int val;
};
struct Recursive {
  static int height(Tree* root){
    if(root){
      auto lh = height(root->left);
      auto rh = height(root->right);
      return 1 + std::max(lh,rh);
    }else{
      return 0;
   }
}
};

This code builds up the stack which is a scarce resource compared to the heap. An iterative solution is shown below:

struct Iterative {
 static int height(Tree *root){
  enum class item_type { tree, height};
  struct item {
    item_type type;
    union {
      Tree* root;
      unsigned height;
    };
    item(Tree* r):type(item_type::tree)
    {root=r;}
    item(unsigned h):type(item_type::height)
    {height=h;}

  };
    bool go_on=true;
    unsigned depth = 0;
    std::stack<item> items;
    while (go_on) {

      if (root) {
        item it(root->right);
        items.push(it);
        root = root->left;
      } else {
        go_on=false;
        depth=0;
        while (!items.empty()&& !go_on) {
          item cur = items.top();
          items.pop();
          switch (cur.type) {
            case item_type::tree: {
                item next(depth);
                items.push(next);
                root = cur.root;
                go_on = true;
              }
            break;
            case item_type::height: {
                  depth = 1 + std::max(depth, cur.height);
                  //depth = 1 + depth + cur.height;
                }
             break;
          }
        }
      }
    }
    return depth;
  }
};

The rest of the blog documents my journey from an OCaml solution to a C++ version. I started with attempting to have a tail recursive solution in OCaml. This lead me to CPS (continuation passing style). The CPS version was then defunctionalized to two mutually recursive functions that was then translated to C++ and then inlined to the iterative method shown above.

Continuation Passing Style
The recursive OCaml solution is shown below:

type 'a tree = Leaf of a' | Node of a' * a' tree * a' tree;;
let depth = function
  | Leaf x -> 0
  | Node(_, left,right) -> 
         let l = depth left in
         let r = depth right in
         1 + (max l r)
                           

The CPS modification as Gasche suggests is straightforward as shown below:

let depth tree =
  let rec depth tree k = match tree with
    | Leaf x -> k 0
    | Node(_,left,right) ->
      depth left (fun dleft ->
        depth right (fun dright ->
          k (1 + (max dleft dright))))
  in depth tree (fun d -> d)

Gasche then defunctionalizes the above code as follows:

type ('a, 'b) cont =
  | Kleft of 'a tree * ('a, 'b) cont (* right and k *)
  | Kright of 'b * ('a, 'b) cont     (* dleft and k *)
  | Kid

let depth tree =
  let rec depth tree k = match tree with
    | Leaf x -> eval k 0
    | Node(_,left,right) ->
      depth left (Kleft(right, k))
  and eval k d = match k with
    | Kleft(right, k) ->
      depth right (Kright(d, k))
    | Kright(dleft, k) ->
      eval k (1 + max d dleft)
    | Kid -> d
  in depth tree Kid
;;

As Gasche points out cont is essentially a linked list and hence the above code can be transformed as shown below:

type ('a, 'b) next_item =
  | Kleft of 'a tree
  | Kright of 'b

type ('a, 'b) cont = ('a, 'b) next_item list

let depth tree =
  let rec depth tree k = match tree with
    | Leaf x -> eval k 0
    | Node(_,left,right) ->
      depth left (Kleft(right) :: k)
  and eval k d = match k with
    | Kleft(right) :: k ->
      depth right (Kright(d) :: k)
    | Kright(dleft) :: k ->
      eval k (1 + max d dleft)
    | [] -> d
  in depth tree []
;;

However the list is accessed like a stack with adding and removing only from the top. From here the C++ solution quickly follows.

type ('a, 'b) next_item =
  | Kleft of 'a tree
  | Kright of 'b

can be translated to C++ as:

enum class item_type { tree, height}; 
struct item {
item_type type;
union { Tree* root; unsigned height; };
};

The constructors for item are just syntactic sugar. The rest of the OCaml code can be translated as:

struct Iterative_1 {

  static int height_aux(Tree *root, stack<tem>& items){
    if(root)
    {
      item it(root->right);
      items.push(it);
      return height_aux(root->left,items);
    } else {
      return eval(items, 0);
    }

  }
  static int eval(stack<item>& items, unsigned depth){
    if(items.empty()){
      return depth;
    }else{

      item cur = items.top();
      items.pop();
      switch (cur.type)
      {
        case item_type::tree:
          {
            item next(depth);
            items.push(next);
            return height_aux(cur.root, items);
          }
        case item_type::height:
          {
            return eval(items,  1 + std::max(depth, cur.height));
          }
      }
    }

  }
static unsigned height(Tree* root){
    std::stack<item>items;
    return height_aux(root,items);
}

Since the recursive calls are all tail recursive we could replace each recursive call with goto. But goto in code is bad for your health and that of the code. Hence we replace it with loops as shown below.

 


static int height_aux(Tree *root, std::stack<item>& items){
    while (root)
    {
      item it(root->right);
      items.push(it);
      root= root->left;
    }
      return eval(items, 0);

  }
  static int eval(std::stack<item>&  items, unsigned depth){
    while (!items.empty()){

      item cur = items.top();
      items.pop();
      switch (cur.type)
      {
        case item_type::tree:
          {
            item next(depth);
            items.push(next);
            return height_aux(cur.root, items);
          }
        case item_type::height:
          {
            depth =  1 + std::max(depth, cur.height);
          }
      }
    }

      return depth;
  }

Now if we inline eval inside height_aux a direct translation we would still require a goto.
The variable go_on as shown in the introduction is a work-around for goto

Additional Computation
Modifying the the iterative to compute additional properties such the number of nodes is done by replacing

depth = 1 + std::max(depth, cur.height);

with

depth = 1 + depth + cur.height;

Computing values such as the maximum value of all the nodes, or its mean or variance should now be trivial.

Personal Note
I would not have come up this solution but for the OCaml code. Thus learning a functional programming language is of great use even for writing code in languages like C++/C#/Java.

About The Sunday Programmer

Joe is an experienced C++/C# developer on Windows. Currently looking out for an opening in C/C++ on Windows or Linux.
This entry was posted in Uncategorized. Bookmark the permalink.

Leave a Reply

Fill in your details below or click an icon to log in:

WordPress.com Logo

You are commenting using your WordPress.com account. Log Out /  Change )

Google photo

You are commenting using your Google account. Log Out /  Change )

Twitter picture

You are commenting using your Twitter account. Log Out /  Change )

Facebook photo

You are commenting using your Facebook account. Log Out /  Change )

Connecting to %s