Online Judge Solutions

Sunday, December 14, 2014

Implement Iterator of Binary Search Tree

Design an iterator over a binary search tree with the following properties:
  1. Elements are visited in ascending order (i.e. an inorder traversal)
  2. next() and hasNext() queries run in O(1) time in average.
Example
For the following binary search tree, inorder traversal by using iterator is [1, 6, 10, 11, 12]
      10
    /     \
 1          11
    \           \
       6           12

 
/**
 * Definition of TreeNode:
 * class TreeNode {
 * public:
 *     int val;
 *     TreeNode *left, *right;
 *     TreeNode(int val) {
 *         this->val = val;
 *         this->left = this->right = NULL;
 *     }
 * }
 * Example of iterate a tree:
 * Solution iterator = Solution(root);
 * while (iterator.hasNext()) {
 *    TreeNode * node = iterator.next();
 *    do something for node
 */
class Solution {
private:
    stack<TreeNode *> st;
    
    void pushLefts(TreeNode *root){
        while(root) {
            st.push(root);
            root = root->left;
        }
    }
    
public:
    //@param root: The root of binary tree.
    Solution(TreeNode *root) {
        pushLefts(root);
    }

    //@return: True if there has next node, or false
    bool hasNext() {
        return st.size() > 0;
    }
    
    //@return: return next node
    TreeNode* next() {
        TreeNode *top = st.top();
        st.pop();
        pushLefts(top->right);
        return top;
    }
};

No comments:

Post a Comment