173. Binary Search Tree Iterator

2024年6月9日 189次阅读

问题

实现一个顺序访问有序二叉树的迭代器。

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.

Example:

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

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.

代码
Compare this typical iterative inorder traversal

  TreeNode visit = root;
  Stack<TreeNode> stack = new Stack();
  while (visit != null || !stack.empty()) {
     while (visit != null) {
          stack.push(visit);
          visit = visit.left;
      }
      TreeNode next = stack.pop();
      visit = next.right;
      doSomethingWith(next.val);
  }

with what we’re supposed to support here:

  BSTIterator i = new BSTIterator(root);
  while (i.hasNext())
      doSomethingWith(i.next());

You can see they already have the exact same structure:

Some initialization.
A while-loop with a condition that tells whether there is more.
The loop body gets the next value and does something with it.
So simply put the three parts of that iterative solution into our three iterator methods:

public class BSTIterator {

    private TreeNode visit;
    // 空间复杂度是O(h)
    private Stack<TreeNode> stack;
    
    public BSTIterator(TreeNode root) {
        visit = root;
        stack = new Stack();
    }
   //时间复杂度是O(1)
    public boolean hasNext() {
        return visit != null || !stack.empty();
    }
    //时间复杂度是O(h)
    public int next() {
        while (visit != null) {
            stack.push(visit);
            visit = visit.left;
        }
        TreeNode next = stack.pop();
        visit = next.right;
        return next.val;
    }
}
点赞