Binary Search Tree Iterator

  • medium
  • tag: BST, stack
  • similar: iterator, iteration(pre/in/post) traversal
  • 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.

思路

  • 一开始看到题目就总是想到in-order traverse得到array, 然后一个个输出即可, 当然不能那么简单, 碎单time可以, 但是space则是O(n)了, 不符合要求. 但也提醒了我, 正如之前做的: Kth Smallest Element in a BST. 可以动态的实现in-order就可以解决了.
  • 首先看到这个O(1), 首先就想到2个数据结构: stack, Map(Array). 前者是按顺序, 直接pop. 后者是因为通过hash拿数据. 这里应该是stack.
  • 既然是stack, 而且next() return 下一个最小的数. 所以stack建立的时候是从大到小的push. 当然利用BST的性质.

idea 0

  • 我先是画图, 先看了随便一个节点的next, 有可能是右子树的做孙子, 也可能是他的root, 也有可能是最上面的root node, 例如root的左子树的右孙子, 那么next应该是root node了.
  • 想着好几个不同的情况: 若node有右子树, 那么next应该返回的是右子树的左叶子节点. 然后返回了该叶子节点后, next应该返回他的root他的并没有返回的root. 不过没有parent pointer, 怎么得到他的root?
  • 总之, 我想找的是规律. 找到了吗?

idea 1

  • stack保存一个root的左子孙, 即所有左儿子的左儿子的左儿子...;
  • 当pop出来的时候, 同时push这个tmp node的右儿子的所有左儿子的左儿子的左儿子...; 其实就是in-order.
  • 参考的 siyang2
  • 代码如下
  public BSTIterator(TreeNode root) {
    stack = new Stack<>();
    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 tmp = stack.pop();
    pushAll(tmp.right);
    return tmp.val;
  }

  private void pushAll(TreeNode root) {
    while (root != null) {
      stack.push(root);
      root = root.left;
    }
  }