LeetCode230-二叉搜索树中第K小的元素

发表于 | 分类于

题目链接

英文链接：https://leetcode.com/problems/kth-smallest-element-in-a-bst/

中文链接：https://leetcode-cn.com/problems/kth-smallest-element-in-a-bst/

题目详述

给定一个二叉搜索树，编写一个函数 kthSmallest 来查找其中第 k 个最小的元素。

说明：
你可以假设 k 总是有效的，1 ≤ k ≤ 二叉搜索树元素个数。

示例 1:

1
2
3
4
5
6
7
输入: root = [3,1,4,null,2], k = 1
   3
  / \
 1   4
  \
   2
输出: 1

示例 2:

1
2
3
4
5
6
7
8
9
输入: root = [5,3,6,2,4,null,null,1], k = 3
       5
      / \
     3   6
    / \
   2   4
  /
 1
输出: 3

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

题目详解

  • 因为是一棵二叉搜索树，所以它的中序遍历序列是有序的，很容易得到第 K 小的元素。
  • 中序遍历可以采用递归或迭代实现。

方法一：递归。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
public class LeetCode_00230 {

    public int kthSmallest(TreeNode root, int k) {
        int[] res = new int[] {0, 0};
        inOrder(root, k, res);
        return res[1];
    }

    private void inOrder(TreeNode root, int k, int[] res) {
        if (root == null) {
            return;
        }
        inOrder(root.left, k, res);
        if (++res[0] == k) {
            res[1] = root.val;
            return;
        }
        inOrder(root.right, k, res);
    }
}

方法二：迭代。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
public class LeetCode_00230 {

    public int kthSmallest(TreeNode root, int k) {
        Stack<TreeNode> stack = new Stack<>();
        while (!stack.isEmpty() || root != null) {
            while (root != null) {
                stack.push(root);
                root = root.left;
            }
            root = stack.pop();
            if (--k == 0) {
                return root.val;
            }
            root = root.right;
        }
        return 0;
    }
}

我们可以转换一下思路，我们已经知道 LeetCode215-数组中的第K个最大元素 的解决方法，本题与它是类似的，只不过对于二叉树而言需要一个统计某棵树的结点个数的方法（在数组中可以通过索引得到个数）。

  • 实现统计结点个数的方法。
  • 若根结点的左子树的结点个数大于 K - 1，说明第 K 小元素在它的左子树中，递归调用即可。
  • 若根结点的左子树的结点个数小于 K - 1，说明第 K 小元素在它的右子树中，递归调用即可，注意参数变化。
  • 若根结点的左子树的结点个数等于 K - 1，说明根结点就是第 K 小的元素，直接返回。
1
2
3
4
5
6
7
8
9
10
11
12
13
public class LeetCode_00230 {

    public int kthSmallest(TreeNode root, int k) {
        int leftSize = treeSize(root.left);
        if (leftSize > k - 1) return kthSmallest(root.left, k);
        if (leftSize < k - 1) return kthSmallest(root.right, k - leftSize - 1);
        return root.val;
    }

    private int treeSize(TreeNode root) {
        return root == null ? 0 : 1 + treeSize(root.left) + treeSize(root.right);
    }
}

对于进阶问题，我们可以修改 TreeNode 结点定义，添加一个 size 变量，在插入和删除结点时，维护路径上结点的 size，使频繁地查找第 k 小的值是比较高效的。