leetcode 173. Binary Search Tree Iterator

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.

代码模板如下：

public class BSTIterator {    public BSTIterator(TreeNode root) {            }    /** @return whether we have a next smallest number */    public boolean hasNext() {            }    /** @return the next smallest number */    public int next() {            }}/** * Your BSTIterator will be called like this: * BSTIterator i = new BSTIterator(root); * while (i.hasNext()) v[f()] = i.next(); */

其实就是中根遍历BST树，把元素以从小到大的顺序遍历出来。

package leetcode;import java.util.Stack;public class BSTIterator {Stack<TreeNode> stack;public BSTIterator(TreeNode root) {stack=new Stack<TreeNode>();while(root!=null){stack.push(root);root=root.left;}}/** @return whether we have a next smallest number */public boolean hasNext() {return !stack.isEmpty();}/** @return the next smallest number */public int next() {TreeNode node=stack.pop();int val=node.val;if(node.right!=null){stack.push(node.right);node=node.right;while(node.left!=null){stack.push(node.left);node=node.left;}}return val;}public static void main(String[] args) {TreeNode root=new TreeNode(4);root.left=new TreeNode(1);root.left.right=new TreeNode(3);root.left.right.left=new TreeNode(2);root.right=new TreeNode(5);root.right.right=new TreeNode(7);root.right.right.left=new TreeNode(6);BSTIterator iterator=new BSTIterator(root);while(iterator.hasNext()){System.out.print(iterator.next()+" ");}}}

我的test用例的图示：

大神思路跟我一样的：

使用stack来存储root的直接左子结点，当next()被调用时，pop出一个元素，并且把这个元素的右结点当作新的root。

public class BSTIterator {    private Stack<TreeNode> stack = new Stack<TreeNode>();        public BSTIterator(TreeNode root) {        pushAll(root);    }    /** @return whether we have a next smallest number */    public boolean hasNext() {        return !stack.isEmpty();    }    /** @return the next smallest number */    public int next() {        TreeNode tmpNode = stack.pop();        pushAll(tmpNode.right);        return tmpNode.val;    }        private void pushAll(TreeNode node) {        for (; node != null; stack.push(node), node = node.left);    }}

这能满足空间复杂度 O(h)。hasNext()的时间复杂度为 O(1) 。next()的平均时间复杂度也是O(1)，因为next()最多被调用n次，而该方法中的 self.pushAll(tmpNode.right) 函数中的右结点的最大个数为n（n为树的结点总数）， 所以平均下来时间复杂度是O(1)。