LwwtCode 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.

调用 next() 将返回二叉搜索树中的下一个最小的数。

Calling next() will return the next smallest number in the BST.

示例：

img

BSTIterator iterator = new BSTIterator(root);
iterator.next();    // 返回 3
iterator.next();    // 返回 7
iterator.hasNext(); // 返回 true
iterator.next();    // 返回 9
iterator.hasNext(); // 返回 true
iterator.next();    // 返回 15
iterator.hasNext(); // 返回 true
iterator.next();    // 返回 20
iterator.hasNext(); // 返回 false

提示：

• next() 和 hasNext() 操作的时间复杂度是 O(1)，并使用 O(h) 内存，其中 h 是树的高度。

• 你可以假设 next() 调用总是有效的，也就是说，当调用 next() 时，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.

• You may assume that next() call will always be valid, that is, there will be at least a next smallest number in the BST when next() is called.

解题思路：

搜索二叉树的中序遍历结果应当是从小到大排序的。所以最简单的方法是中序遍历二叉树把结果存入数组，判断是否为从小到大排序：

list = Inorder(root)
return list == sort(list)

但是这种把所有结点值都存入数组的空间复杂度为 O(n)，其中 n 为结点数，而不是 二叉树的深度 O(logn)

另一种方法是维护一个栈，递归入栈二叉树的左子结点。每次调用 next() 函数时，出栈一个结点，返回其结点值。并且如果该结点存在右子结点，则递归入栈该结点左子树的左子结点，准备下一次操作。空间复杂度为栈的容量，最大为二叉树的最大深度。

Java

class BSTIterator {
    Stack<TreeNode> stack;

    public BSTIterator(TreeNode root) {
        this.stack = new Stack<>();
        leftInorder(root);
    }

    private void leftInorder(TreeNode node) {
        while (node != null) {
            this.stack.add(node);
            node = node.left;
        }
    }

    /** @return the next smallest number */
    public int next() {
        TreeNode node = this.stack.pop();
        if (node.right != null)
            leftInorder(node.right);
        return node.val;
    }

    /** @return whether we have a next smallest number */
    public boolean hasNext() {
        return this.stack.size() > 0;
    }
}

/**
 * Your BSTIterator object will be instantiated and called as such:
 * obj := Constructor(root);
 * param_1 := obj.Next();
 * param_2 := obj.HasNext();
 */

Python

class BSTIterator:

    def __init__(self, root: TreeNode):
        self.stack = []
        self._left_inorder(root)

    def _left_inorder(self, node: TreeNode):
        while node:
            self.stack.append(node)
            node = node.left

    def next(self) -> int:
        """
        @return the next smallest number
        """
        node = self.stack.pop()
        if self.stack.right:
            self._left_inorder(self.stack.right)
        return node.val

    def hasNext(self) -> bool:
        """
        @return whether we have a next smallest number
        """
        return len(self.stack) > 0

Golang

type BSTIterator struct {
}

// 用切片 slice 实现的栈
var stack []*TreeNode

func Constructor(root *TreeNode) BSTIterator {
    _leftInorder(root)
    return BSTIterator{}
}

func _leftInorder(root *TreeNode) {
    for root != nil {
        stack = append(stack, root)
        root = root.Left
    }
}

/** @return the next smallest number */
func (this *BSTIterator) Next() int {
    root := stack[len(stack)-1]
    stack = stack[:len(stack)-1]

    if root.Right != nil {
        _leftInorder(root.Right)
    }
    return root.Val
}

/** @return whether we have a next smallest number */
func (this *BSTIterator) HasNext() bool {
    return len(stack) > 0
}