【leetcode】Binary Search Tree Iterator（middle）

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) {
        pCur = root;
        while(pCur != NULL)
        {
            v.push_back(pCur);
            pCur = pCur->left;
        }
    }

    /** @return whether we have a next smallest number */
    bool hasNext() {
        return (!v.empty() || NULL != pCur);
    }

    /** @return the next smallest number */
    int next() {
        int num;
        TreeNode * tmp = v.back();
        v.pop_back();
        num = tmp->val;
        pCur = tmp->right;
        while(pCur != NULL)
        {
            v.push_back(pCur);
            pCur = pCur->left;
        }
        return num;
    }
private:
    vector<TreeNode *> v;
    TreeNode * pCur;
};

 

大神更精简的代码： 经验，把相同功能的代码放在一起可以简化代码。

public class BSTIterator {

        Stack<TreeNode> stack =  null ;            
        TreeNode current = null ;

        public BSTIterator(TreeNode root) {
              current = root;        
              stack = new Stack<> ();
        }

        /** @return whether we have a next smallest number */
        public boolean hasNext() {        
              return !stack.isEmpty() || current != null;  
        }

            /** @return the next smallest number */
        public int next() {
            while (current != null) {
                stack.push(current);
                current = current.left ;
            }       
            TreeNode t = stack.pop() ;      
            current = t.right ;     
            return t.val ;
        }
    }