Binary tree iterator algorithm

Binary tree pre-order,in-order and post-order traversal are basics in algorithm and data structure.And the recursive binary tree traversal is a classical application of recursion.

We define a binary tree node as :

// C++
struct Node {
    int value;
    Node *left;
    Node *right;
}

In order binary tree traversal can be:

// C++
void inorder_traverse(Node *node) {
    if (NULL != node->left) {
        inorder_traverse(node->left);
    }
    do_something(node);
    if (NULL != node->right) {
        inorder_traverse(node->right);
    }
}

Easy enough, right? Pre-order and post-order traversal algorithms are similar.

In many applications, we may want to abstract the traversal itself. Now assume we have a function to add all nodes in the tree, we want to use it on linked list, array or binary tree. The best way is to have an iterator so that it can decouple the algorithm and data structure, so that an algorithm can be used on different data structures. We define sum as :

int sum(Iterator it)

Linked list is a linear structure, so its iterator implementation is simple and straightforward. While iterator of binary tree is not so easy to implement, the reason is the binary tree traversal is automatically performed by the compiler on the call stack, programmers usually don't have enough control over this process. Can we control this process by ourselves? Yes, we can use a non-recursive method:

// C++
void inorder_traverse_nonrecursive(Node *node) {
    Stack stack;
    do {
        while (NULL != node) {
            stack.push(node);
            node = node->left;
        }
        do {
            Node *top = stack.top();
            stack.pop();
            do_something(top);
            if (NULL != top->right) {
                node = top->right;
                break;
            }
        }
        while (!stack.empty());
    }
    while (!stack.empty());
}

Now we have control over the traversal. Next we will implement the iterator. The key here is the traversal process depends on the status of the stack, so we need to have a stack.The interface of the iterator is:

// C++
class Iterator {
    public:
        virtual Node* next() = 0;
};

Next is the implementation:

//C++
class InorderIterator : public Iterator {
    public:
        InorderIterator(Node *node) {
            Node *current = node;
            while (NULL != current) {
                mStack.push(current);
                current = current->left;
            }
        }
        virtual Node* next() {
            if (mStack.empty()) {
                return NULL;
            }
            Node *top = mStack.top();
            mStack.pop();
            if (NULL != top->right) {
                Node *current = top->right;
                while (NULL != current) {
                    mStack.push(current);
                    current = current->left;
                }
            }
            return top;
         }
    private:
        std::stack<Node*> mStack;
};

Let's consider the time and space complexity of this algorithm. Apparently the stack only needs to store n nodes at most, so the space complexity is O(n). How about time complexity? Each next() call needs at most n stack operations as well, and the whole binary tree traversal needs n next() calls, is the time complexity O(n*2) then? No, since each node can be pushed and popped only once, the time complexity is O(n) as well.

Bonus point, if the language supports yield such as C# or Python, we can write the algorithm as:

// Python
def inorder(t):
    if t:
        for x in inorder(t.left):
            yield x
        yield t.label
        for x in inorder(t.right):
            yield x

The difference between yield and return is yield will return the system call status, next time when this function is called again, it will start from the point where it stops last time.

Author : @文艺复兴记 Source : http://coolshell.cn/articles/9886.html#jtss-tsina