Leetcode-173. Binary Search Tree Iterator

题目

Implement an iterator over a binary search tree (BST). Your iterator will be initialized with the root node of a BST.

Calling next() will return the next smallest number in the BST.

Note: next() and hasNext() should run in average O(1) time and uses O(h) memory, where h is the height of the tree.

思路

这道题本质上是写一个二叉树的中序遍历的迭代器。内部设置一个栈，初始化的时候，存储从根节点到最左叶子节点的路径。在遍历的过程中，每次从栈中弹出一个元素，作为当前的返回结果，同时探测一下当前节点是否存在右孩子，如果有，则进入右孩子，并把从该右孩子到最左叶子节点的所有节点入栈。

代码

class BSTIterator {
public:
    BSTIterator(TreeNode *root) {
        while(root) {
            ss.push(root);
            root = root->left;
        }
    }

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

    int next() {
        TreeNode *t = ss.top();
        ss.pop();
        if(t->right) {
            TreeNode *p = t->right;
            while(p) {
                ss.push(p);
                p = p->left;
            }
        }
        return t->val;
    }
private:
    stack<TreeNode*> ss;
};