173. Binary Search Tree Iterator

173. Binary Search Tree Iterator

Description

mplement 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.

Implementation

Data structure is combined with Stack and a TreeNode, the tree node has left and right nodes, and the stack is used to store the next number for smallest number.

From the problem we can know, we can know that the this is a binary search problem with in-order search. So in the hasNode I add the left-part value into it and in the output part with pop node, we will set the right node for current node to search the right part.

/**
 * Definition for binary tree
 * struct TreeNode {
 *     int val;
 *     TreeNode *left;
 *     TreeNode *right;
 *     TreeNode(int x) : val(x), left(NULL), right(NULL) {}
 * };
 */
class BSTIterator {
public:
    TreeNode* cur;
    stack<TreeNode*> stk;

    BSTIterator(TreeNode *root) {
        this->cur = root;
    }

    void addNode(TreeNode* cur) {
        while(this->cur != NULL) {
            stk.push(this->cur);
            this->cur = this->cur->left;
        }
    }

    /** @return whether we have a next smallest number */
    bool hasNext() {
        if(this->cur != NULL) {
            addNode(this->cur);
            return true;
        }
        else if(stk.size() > 0) 
            return true;
        return false;    
    }

    /** @return the next smallest number */
    int next() {
        TreeNode* tp = stk.top();
        stk.pop();
        cur = tp->right;
        return tp->val;
    }
};