230. 二叉搜索树中第K小的元素

题目描述

给定一个二叉搜索树的根节点 root ，和一个整数 k ，请你设计一个算法查找其中第 k 个最小元素（从 1 开始计数）。

示例 1：

输入：root = [3,1,4,null,2]，k = 1

输出：1

示例 2：

输入：root = [5,3,6,2,4,null,null,1]，k = 3

输出：3

提示：

• 树中的节点数为 n 。
• 1 <= k <= n <= 104
• 0 <= Node.val <= 104

进阶：如果二叉搜索树经常被修改（插入/删除操作）并且你需要频繁地查找第 k 小的值，你将如何优化算法？

出处：LeetCode

解法一：中序遍历

思路

中序遍历二叉搜索树，得到一个递增的序列，第 k 个元素即为所求。

实现

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

复杂度分析

• 时间复杂度：$O(n)$
• 空间复杂度：$O(n)$

解法二：递归

思路

递归遍历二叉搜索树，记录当前节点的排名，当排名等于 k 时返回当前节点的值。

实现

function kthSmallest(root: TreeNode | null, k: number): number {
  let res = 0;
  const dfs = (node: TreeNode | null): boolean => {
    if (node === null) {
      return false;
    }
    if (dfs(node.left)) {
      return true;
    }
    if (--k === 0) {
      res = node.val;
      return true;
    }
    return dfs(node.right);
  };
  dfs(root);
  return res;
}

复杂度分析

• 时间复杂度：$O(n)$
• 空间复杂度：$O(n)$

进阶解法：Morris 遍历

思路

Morris 遍历是一种遍历二叉树的方法，可以在 $O(1)$ 的空间复杂度下完成中序遍历。具体做法是：

1. 初始化当前节点 cur 为根节点 root。
2. 如果 cur 为空，结束遍历。
3. 如果 cur 的左子树为空，输出 cur 的值，然后将 cur 更新为 cur 的右子树。
4. 如果 cur 的左子树不为空，找到 cur 的前驱节点 pre。
   • 如果 pre 的右子树为空，将 pre 的右子树指向 cur，然后将 cur 更新为 cur 的左子树。
   • 如果 pre 的右子树不为空，将 pre 的右子树恢复为空，输出 cur 的值，然后将 cur 更新为 cur 的右子树。
5. 重复步骤 2-4，直到 cur 为空。
6. 结束遍历。

实现

function kthSmallest(root: TreeNode | null, k: number): number {
  let res = 0;
  let count = 0;
  let cur: TreeNode | null = root;
  while (cur) {
    if (cur.left === null) {
      count++;
      if (count === k) {
        res = cur.val;
        break;
      }
      cur = cur.right;
    } else {
      let pre = cur.left;
      while (pre.right && pre.right !== cur) {
        pre = pre.right;
      }
      if (pre.right === null) {
        pre.right = cur;
        cur = cur.left;
      } else {
        pre.right = null;
        count++;
        if (count === k) {
          res = cur.val;
          break;
        }
        cur = cur.right;
      }
    }
  }
  return res;
}

复杂度分析

• 时间复杂度：$O(n)$
• 空间复杂度：$O(1)$