【LeetCode】230#二叉搜索树中第K小的元素

题目描述

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

说明：

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

示例 1:

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

示例 2:

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

进阶：

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

---

解题思路

二叉搜索树（BST）是一棵二叉树，每个节点的值都大于其左子树中的任意节点的的值而小于右子树的任意节点的值。

翻阅了一下《算法》第四版的相关章节，发现树上的二叉树维护了一个N来记录以当前节点为根节点的子树的节点总数。通过N，我们就可以从上往下遍历二叉搜索树时，选择往左走还是往右走，直到找到我们要找的节点。具体方式如下：

假设我们想找到第k小的元素（即树中正好有k-1个小于它的值），如果左子树中的节点数t > k-1，那我们就继续（递归地）在左子树中搜索第k小的元素；如果t = k-1我们就返回该节点的值；如果t < k-1，我们就（递归地）在右子树中查找第k-t-1小的值。

现在我们需要一个统计N的方法private int size (TreeNode root)，可以用递归的方式来实现，每个节点的N等于该节点的左子树的N+该节点的右子树的N+1，这个 1 就是该节点。

源代码

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

private int size (TreeNode root) {
    if (root == null) return 0;
    return size(root.left) + size(root.right) + 1;
}

心得体会

这一题是二叉搜索树，而不是简单的二叉树，所以要利用二叉搜索树的特点来找到解题思路。