二叉搜索树迭代器

76 阅读1分钟

持续创作,加速成长!这是我参与「掘金日新计划 · 10 月更文挑战」的第24天,点击查看活动详情

173. 二叉搜索树迭代器 - 力扣(LeetCode)

实现一个二叉搜索树迭代器类BSTIterator ,表示一个按中序遍历二叉搜索树(BST)的迭代器:

  • BSTIterator(TreeNode root) 初始化 BSTIterator 类的一个对象。BST 的根节点 root 会作为构造函数的一部分给出。指针应初始化为一个不存在于 BST 中的数字,且该数字小于 BST 中的任何元素。
  • boolean hasNext() 如果向指针右侧遍历存在数字,则返回 true ;否则返回 false
  • int next()将指针向右移动,然后返回指针处的数字。

注意,指针初始化为一个不存在于 BST 中的数字,所以对 next() 的首次调用将返回 BST 中的最小元素。

你可以假设 next() 调用总是有效的,也就是说,当调用 next() 时,BST 的中序遍历中至少存在一个下一个数字。

示例:

输入
["BSTIterator", "next", "next", "hasNext", "next", "hasNext", "next", "hasNext", "next", "hasNext"]
[[[7, 3, 15, null, null, 9, 20]], [], [], [], [], [], [], [], [], []]
输出
[null, 3, 7, true, 9, true, 15, true, 20, false]
解释
BSTIterator bSTIterator = new BSTIterator([7, 3, 15, null, null, 9, 20]);
bSTIterator.next(); // 返回 3
bSTIterator.next(); // 返回 7
bSTIterator.hasNext(); // 返回 True
bSTIterator.next(); // 返回 9
bSTIterator.hasNext(); // 返回 True
bSTIterator.next(); // 返回 15
bSTIterator.hasNext(); // 返回 True
bSTIterator.next(); // 返回 20
bSTIterator.hasNext(); // 返回 False

提示:

  • 树中节点的数目在范围 [1, 10^5]
  • 0 <= Node.val <= 10^6
  • 最多调用 10^5hasNextnext 操作

进阶:

  • 你可以设计一个满足下述条件的解决方案吗?next()hasNext() 操作均摊时间复杂度为 O(1) ,并使用 O(h) 内存。其中 h 是树的高度。

思路

本题有多种解法。首先想到的就是用一个队列中序遍历所有节点,next()的时候,从队列头部移除元素,hasNext()的时候判断队列是否不为空。

还可以使用yield实现一个迭代器,next()的时候按照中序遍历每次返回一个节点,hasNext()的时候判断迭代器的done

解题

/**
 * Definition for a binary tree node.
 * function TreeNode(val, left, right) {
 *     this.val = (val===undefined ? 0 : val)
 *     this.left = (left===undefined ? null : left)
 *     this.right = (right===undefined ? null : right)
 * }
 */
/**
 * @param {TreeNode} root
 */
var BSTIterator = function (root) {
  this.iterator = this._next(root);
  this.curr = this.iterator.next();
};

/**
 * @return {number}
 */
BSTIterator.prototype.next = function () {
  let val = this.curr.value.val;
  this.curr = this.iterator.next();
  return val;
};

/**
 * @return {boolean}
 */
BSTIterator.prototype.hasNext = function () {
  return !this.curr.done;
};

BSTIterator.prototype._next = function* (node) {
  if (node.left) {
    yield* this._next(node.left);
  }
  yield node;
  if (node.right) {
    yield* this._next(node.right);
  }
};

/**
 * Your BSTIterator object will be instantiated and called as such:
 * var obj = new BSTIterator(root)
 * var param_1 = obj.next()
 * var param_2 = obj.hasNext()
 */

进阶

我们观察中序遍历规则,先遍历最左侧节点,然后再父节点,最后是右节点。可以用一个栈存当前节点的左侧节点,next()的时候,取出栈顶元素,并把该节点右节点及其左子树的左节点入栈,hasNext()的时候判断栈是否不为空。

/**
 * Definition for a binary tree node.
 * function TreeNode(val, left, right) {
 *     this.val = (val===undefined ? 0 : val)
 *     this.left = (left===undefined ? null : left)
 *     this.right = (right===undefined ? null : right)
 * }
 */
/**
 * @param {TreeNode} root
 */
var BSTIterator = function (root) {
  this.stack = [root];
  let node = root;
  while (node.left) {
    this.stack.push(node.left);
    node = node.left;
  }
};

/**
 * @return {number}
 */
BSTIterator.prototype.next = function () {
  let node = this.stack.pop();
  const val = node.val;
  if (node.right) {
    this.stack.push(node.right);
    node = node.right;
    while (node.left) {
      this.stack.push(node.left);
      node = node.left
    }
  }
  return val;
};

/**
 * @return {boolean}
 */
BSTIterator.prototype.hasNext = function () {
  return this.stack.length > 0;
};