題目描述

給定一個二叉搜索樹，編寫一個函數 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;
}

心得體會

這一題是二叉搜索樹，而不是簡單的二叉樹，所以要利用二叉搜索樹的特點來找到解題思路。

推薦閱讀：

相關文章