Binary tree iterators

2018-04-20

Content:

1. Problem statement

Given a binary tree, build iterator classes that can traverse that tree (in preorder, inorder, postorder) and have following operators:

  • next: return current node and move the iterator 1 step.
  • hasNext: check if we’re done traversing the tree or not.

Usage example:

1
2
3
4
5
6

// Inorder binary tree iterator
InorderBTIterator it(root);
while(it.hasNext()) {
  // Visit node in inorder
  print("%d ", it.next()->val);
}

2. Solution

All we need to do is wrap the iterative implementation of the traversal algorithm inside next method, and return the node when we visit it.

3. Implementation

I just write the code for preorder and inorder binary tree iterator, the postorder iterator’s is totally similar.

Preorder binary tree iterator:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24

class PreorderBTIterator {
private:
  stack<TreeNode*> stk;
public:
  PreorderBTIterator(TreeNode* root) {
    if (root != NULL) stk.push(root);
  }

  TreeNode* next() {
    while(!stk.empty()) {
      TreeNode* ret = stk.top();
      stk.pop();

      if (ret->right != NULL) stk.push(ret->right);
      if (ret->left != NULL) stk.push(ret->left);

      return ret;
    }
  }

  bool hasNext() {
    return !stk.empty();
  }
};

Side notes: we can use Morris inorder traversal implementation to reduce the space complexity (basically we can use any iterative tree traversal implementation in next method).

Inorder binary tree iterator:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25

class InorderBTIterator {
private:
  TreeNode* node;
  stack<TreeNode*> stk;
public:		
  InorderBTIterator(TreeNode* root) : node(root) {}

  TreeNode* next() {
    while(!stk.empty() || node != NULL) {
      if (node != NULL) {
        stk.push(node);
        node = node->left;
      } else {
        TreeNode* ret = stk.top();
        stk.pop();
        node = ret->right;
        return ret;
      }
    }
  }

  bool hasNext() {
    return (!stk.empty() || node != NULL);
  }
};