【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 ;

        }

    }