LeetCode 173. Binary Search Tree Iterator

思路

画一个简单的BST，然后可以看出有三种情况

• 如果当前节点有右孩子，那就在右子树里找到值最小的那个，显然是右子树里最左下侧的那个节点
• 如果当前节点没有右孩子，由于左子树的值肯定都比当前节点的值小，左右要找父节点
  • 如果找到某个父节点的值比当前节点值大，那这个父节点就是下一节点
  • 否则，当前节点就是最后一个节点

因此，需要一个栈结构记录当前节点的父亲们（有点像调用栈）。
测试例：NULL根节点，只有一个节点，只有单侧树

代码

/**
 * Definition for binary tree
 * struct TreeNode {
 *     int val;
 *     TreeNode *left;
 *     TreeNode *right;
 *     TreeNode(int x) : val(x), left(NULL), right(NULL) {}
 * };
 */
class BSTIterator {
private:
    stack<TreeNode*> callStack;
    TreeNode* currNode;
    TreeNode* rootNode;
public:
    BSTIterator(TreeNode *root) {
        currNode = NULL;
        rootNode = root;
    }
    /** @return whether we have a next smallest number */
    bool hasNext() {
        if(currNode == NULL){
            currNode = rootNode;
            while(currNode!=NULL && currNode->left != NULL){
                callStack.push(currNode);
                currNode = currNode->left;
            }
            return currNode!=NULL;
        }
        if(currNode->right != NULL){
            //left-most node in right sub-tree
            callStack.push(currNode);
            currNode = currNode->right;
            while(currNode->left != NULL){
                callStack.push(currNode);
                currNode = currNode->left;
            }
            return true;
        } else {
            //look for first parent whos val is > currNode->val. If not available, return false
            int currVal = currNode->val;
            while(!callStack.empty()){
                TreeNode* parent = callStack.top();
                callStack.pop();
                if(parent->val > currVal){
                    currNode = parent;
                    return true;
                }
            }
            currNode = NULL;
            return false;
        }
    }
    /** @return the next smallest number */
    int next() {
        if(currNode != NULL)
            return currNode->val;
        return INT_MIN;
    }
};
/**
 * Your BSTIterator will be called like this:
 * BSTIterator i = BSTIterator(root);
 * while (i.hasNext()) cout << i.next();
 */