Binary Search Tree Iterator

Question

  1. Design an iterator over a binary search tree with the following rules:
  3. - Elements are visited in ascending order (i.e. an in-order traversal)
  4. - next() and hasNext() queries run in O(1) time in average.
  6. Example
  7. For the following binary search tree, in-order traversal by using iterator is [1, 6, 10, 11, 12]
  9.            10
  10.          /    \
  11.         1      11
  12.          \       \
  13.              6       12
  15. Challenge
  16. Extra memory usage O(h), h is the height of the tree.
  18. Super Star: Extra memory usage O(1)

题解 - 中序遍历

仍然考的是中序遍历，但是是非递归实现。其实这道题等价于写一个二叉树中序遍历的迭代器。需要内置一个栈，一开始先存储到最左叶子节点的路径。在遍历的过程中，只要当前节点存在右子树，则进入右子树，存储从此处开始到当前子树里最左叶子节点的路径。

C++

  1. /**
  2.  * Definition of TreeNode:
  3.  * class TreeNode {
  4.  * public:
  5.  *     int val;
  6.  *     TreeNode *left, *right;
  7.  *     TreeNode(int val) {
  8.  *         this->val = val;
  9.  *         this->left = this->right = NULL;
  10.  *     }
  11.  * }
  12.  * Example of iterate a tree:
  13.  * BSTIterator iterator = BSTIterator(root);
  14.  * while (iterator.hasNext()) {
  15.  *    TreeNode * node = iterator.next();
  16.  *    do something for node
  17.  */
  18.     class BSTIterator {
  19.     private:
  20.         stack<TreeNode*> stack_;
  21.         TreeNode* cur_ = NULL;
  23.     public:
  24.         //@param root: The root of binary tree.
  25.         BSTIterator(TreeNode *root) {
  26.             // write your code here
  27.             cur_ = root;
  28.         }
  30.         //@return: True if there has next node, or false
  31.         bool hasNext() {
  32.             // write your code here
  33.             return (cur_ || !stack_.empty());
  34.         }
  36.         //@return: return next node
  37.         TreeNode* next() {
  38.             // write your code here
  39.             while (cur_) {
  40.                 stack_.push(cur_);
  41.                 cur_ = cur_->left;
  42.             }
  43.             cur_ = stack_.top();
  44.             stack_.pop();
  45.             TreeNode* node = cur_;
  46.             cur_ = cur_->right;
  48.             return node;
  49.         }
  50.     };

Java

  1. /**
  2.  * Definition of TreeNode:
  3.  * public class TreeNode {
  4.  *     public int val;
  5.  *     public TreeNode left, right;
  6.  *     public TreeNode(int val) {
  7.  *         this.val = val;
  8.  *         this.left = this.right = null;
  9.  *     }
  10.  * }
  11.  * Example of iterate a tree:
  12.  * Solution iterator = new Solution(root);
  13.  * while (iterator.hasNext()) {
  14.  *    TreeNode node = iterator.next();
  15.  *    do something for node
  16.  * } 
  17.  */
  18. public class Solution {
  19.     private Stack<TreeNode> stack = new Stack<>();
  20.     private TreeNode curt;
  22.     // @param root: The root of binary tree.
  23.     public Solution(TreeNode root) {
  24.         curt = root;
  25.     }
  27.     //@return: True if there has next node, or false
  28.     public boolean hasNext() {
  29.         return (curt != null || !stack.isEmpty()); //important to judge curt != null
  30.     }
  32.     //@return: return next node
  33.     public TreeNode next() {
  34.         while (curt != null) {
  35.             stack.push(curt);
  36.             curt = curt.left;
  37.         }
  39.         curt = stack.pop();
  40.         TreeNode node = curt;
  41.         curt = curt.right;
  43.         return node;
  44.     }
  45. }

源码分析

  1. 这里容易出错的是 hasNext() 函数中的判断语句，不能漏掉 curt != null。
  2. 如果是 leetcode 上的原题，由于接口不同，则不需要维护 current 指针。