[LintCode] Binary Search Tree Iterator

最新推荐文章于 2020-05-01 14:07:25 发布

转载最新推荐文章于 2020-05-01 14:07:25 发布 · 76 阅读

0 ·

CC 4.0 BY-SA版权

原文链接：https://segmentfault.com/a/1190000005032130

本文介绍了一种高效的二叉搜索树迭代器设计方法，该方法能够在平均O(1)的时间复杂度内完成next()和hasNext()操作。通过使用栈结构预先加载左子树节点，确保了元素按升序遍历。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

Problem

Design an iterator over a binary search tree with the following rules:

Elements are visited in ascending order (i.e. an in-order traversal)
next() and hasNext() queries run in O(1) time in average.

Example

For the following binary search tree, in-order traversal by using iterator is [1, 6, 10, 11, 12]

Challenge

Extra memory usage O(h), h is the height of the tree.

Super Star: Extra memory usage O(1)

Note

建立一个堆栈，先将最左边的结点从大到小压入栈，这样的话，为了实现迭代即返回下一个node的next()函数就要考虑右边的结点。比如example中的BST，stack第一次pop出来的结点是1，然后就应该把它的右子树结点6压入stack；又如pop出了root结点10以后，就要把11压入堆栈，这样，在pop出11之后，再将12压入堆栈。如此，实现next()函数。

Solution

public class BSTIterator {
    //@param root: The root of binary tree.
    Stack<TreeNode> stack;
    public BSTIterator(TreeNode root) {
        stack = new Stack<TreeNode>();
        while (root != null) {
            stack.push(root);
            root = root.left;
        }
    }

    //@return: True if there has next node, or false
    public boolean hasNext() {
        return !stack.isEmpty();
    }
    
    //@return: return next node
    public TreeNode next() {
        TreeNode cur = stack.pop();
        TreeNode temp = cur;
        if (cur.right != null) {
            cur = cur.right;
            while (cur != null) {
                stack.push(cur);
                cur = cur.left;
            }
        }
        return temp;
    }
}