230. Kth Smallest Element in a BST

Leetcode Tree Binary Search

Given a binary search tree, write a function kthSmallest to find the kth smallest element in it.

Note: You may assume $k$ is always valid, $1 ≤ k ≤$ BST's total elements.

Example 1:

Input: root = [3,1,4,null,2], k = 1
   3
  / \
 1   4
  \
   2
Output: 1

Example 2:

Input: root = [5,3,6,2,4,null,null,1], k = 3
       5
      / \
     3   6
    / \
   2   4
  /
 1
Output: 3

Follow up:

• What if the BST is modified (insert/delete operations) often and you need to find the $k$th smallest frequently? How would you optimize the $k$thSmallest routine?

分析

这道题目要求我们求出第$k$个小的节点。既然是二叉搜索树，由于二叉搜索树的中序遍历(94. Binary Tree Inorder Traversal )的结果是递增序列。所以最直接的方法就是进行二叉搜索树的中序遍历，当遍历到第$k$个节点时，返回该节点。

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

但这种方案是比较低效的，万一$k$值接近于节点个数呢？也就说这种算法在最坏情况下的时间复杂度是$O(n)$，而且由于要保存节点，在最坏情况下空间复杂度也是$O(n)$。有没有更快的方法呢？当然是有的。二叉搜索树支持有序的方法(order-based methods)，它支持计算子树的节点的个数。一开始可以计算左子树的个数，当左子树节点的个数小于$k$时，说明第$k$个节点必定在左子树；当等于左子树节点的个数加1时，说明根节点就是第$k$个小的节点；否则寻找右子树。该方法非常类似与floor(x):返回小于或者等于$x$的最大节点。

```Java public int kthSmallest(TreeNode root, int k) { int count = size(root.left); if (k <= count) return kthSmallest(root.left, k); else if (k > count + 1) return kthSmallest(root.right, k - 1 - count); // 1 is counted as current node return root.val; }

// 树的节点个数 public int size(TreeNode root) { if (root == null) return 0; return 1 + size(root.left) + size(root.right); } ```