第K小与第K大元素

为什么中序遍历能定位第 k 小？

• 依据: 中序遍历按升序输出, 第 k 个出栈的节点就是第 k 小;
• 优化: 不必完整遍历, 计数到 k 就返回;

const kthSmallest = function (root, k) {
  const stack = [];
  let node = root;
  while (node != null || stack.length !== 0) {
    while (node != null) {
      stack.push(node);
      node = node.left;
    }
    node = stack.pop();
    k--;
    if (k === 0) return node.val;
    node = node.right;
  }
};

第 k 大如何复用同一套代码？

• 做法: 交换左右遍历顺序, 改成 "右 - 根 - 左", 出栈顺序变成降序;
• 结果: 第 k 个出栈的元素就是第 k 大;

const findTargetNode = function (root, k) {
  const stack = [];
  let node = root;
  while (node != null || stack.length !== 0) {
    while (node != null) {
      stack.push(node);
      node = node.right;
    }
    node = stack.pop();
    k--;
    if (k === 0) return node.val;
    node = node.left;
  }
};

复杂度如何？

• 时间: O(h + k), h 为树高; 找第 1 大只需下降到最右侧;
• 空间: O(h), 即栈内未回溯的祖先节点;
• 关联: 迭代遍历的通用写法见 遍历的迭代实现;