LeetCode

Binary Search Tree Iterator

2015.1.23 18:58

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.

Solution:

　　This is a BST, so the next smallest element is exactly the inorder successor. Make a whole inorder traversal and you'll get every successor.

　　Use a hash map to record the result, so that you can retrieve them in O(1) time later.

　　The traversal will be done immediately after receiving the root, that is, inside the constructor.

　　Time complexity on average is O(1), space complexity is O(n).

Accepted code:

// 1CE, 1MLE, 1WA, 1AC, anonther good problem
#include <unordered_map>
using namespace std;
/**
 * Definition for binary tree
 * struct TreeNode {
 *     int val;
 *     TreeNode *left;
 *     TreeNode *right;
 *     TreeNode(int x) : val(x), left(NULL), right(NULL) {}
 * };
 */
class BSTIterator {
public:
    BSTIterator(TreeNode *root) {
        nn.clear();
        currentNode = nullptr;
        
        if (root == nullptr) {
            return;
        }
        
        TreeNode *ll, *rr;
        findSuccessor(root, ll, rr);
        currentNode = ll;
    }

    /** @return whether we have a next smallest number */
    bool hasNext() {
        return currentNode != nullptr;
    }

    /** @return the next smallest number */
    int next() {
        int val = currentNode->val;
        currentNode = nn[currentNode];
        return val;
    }
    
    ~BSTIterator() {
        nn.clear();
    }
private:
    unordered_map<TreeNode *, TreeNode *> nn;
    TreeNode *currentNode;
    
    void findSuccessor(TreeNode *root, TreeNode *&left, TreeNode *&right) {
        TreeNode *ll, *lr, *rl, *rr;
        
        if (root->left == nullptr) {
            ll = root;
            lr = root;
        } else {
            findSuccessor(root->left, ll, lr);
            nn[lr] = root;
        }
        
        if (root->right == nullptr) {
            rl = root;
            rr = root;
        } else {
            findSuccessor(root->right, rl, rr);
            nn[root] = rl;
        }
        left = ll;
        right = rr;
    }
};

/**
 * Your BSTIterator will be called like this:
 * BSTIterator i = BSTIterator(root);
 * while (i.hasNext()) cout << i.next();
 */