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 kth smallest frequently? How would you optimize the kthSmallest routine?
Credits:
Special thanks to @ts for adding this problem and creating all test cases.
這又是一道關於二叉搜索樹 Binary Search Tree 的題， LeetCode 中關於 BST 的題有 Validate Binary Search Tree， Recover Binary Search Tree， Binary Search Tree Iterator， Unique Binary Search Trees， Unique Binary Search Trees II，Convert Sorted Array to Binary Search Tree 和 Convert Sorted List to Binary Search Tree。那么這道題給的提示是讓我們用 BST 的性質來解題，最重要的性質是就是左<根<右，如果用中序遍歷所有的節點就會得到一個有序數組。所以解題的關鍵還是中序遍歷啊。關於二叉樹的中序遍歷可以參見我之前的博客 Binary Tree Inorder Traversal，里面有很多種方法可以用，先來看一種非遞歸的方法，中序遍歷最先遍歷到的是最小的結點，只要用一個計數器，每遍歷一個結點，計數器自增1，當計數器到達k時，返回當前結點值即可，參見代碼如下：
解法一：
class Solution { public: int kthSmallest(TreeNode* root, int k) { int cnt = 0; stack<TreeNode*> s; TreeNode *p = root; while (p || !s.empty()) { while (p) { s.push(p); p = p->left; } p = s.top(); s.pop(); ++cnt; if (cnt == k) return p->val; p = p->right; } return 0; } };
當然，此題我們也可以用遞歸來解，還是利用中序遍歷來解，代碼如下：
解法二：
class Solution { public: int kthSmallest(TreeNode* root, int k) { return kthSmallestDFS(root, k); } int kthSmallestDFS(TreeNode* root, int &k) { if (!root) return -1; int val = kthSmallestDFS(root->left, k); if (k == 0) return val; if (--k == 0) return root->val; return kthSmallestDFS(root->right, k); } };
再來看一種分治法的思路，由於 BST 的性質，可以快速定位出第k小的元素是在左子樹還是右子樹，首先計算出左子樹的結點個數總和 cnt，如果k小於等於左子樹結點總和 cnt，說明第k小的元素在左子樹中，直接對左子結點調用遞歸即可。如果k大於 cnt+1，說明目標值在右子樹中，對右子結點調用遞歸函數，注意此時的k應為 k-cnt-1，應為已經減少了 cnt+1 個結點。如果k正好等於 cnt+1，說明當前結點即為所求，返回當前結點值即可，參見代碼如下：
解法三：
class Solution { public: int kthSmallest(TreeNode* root, int k) { int cnt = count(root->left); if (k <= cnt) { return kthSmallest(root->left, k); } else if (k > cnt + 1) { return kthSmallest(root->right, k - cnt - 1); } return root->val; } int count(TreeNode* node) { if (!node) return 0; return 1 + count(node->left) + count(node->right); } };
這道題的 Follow up 中說假設該 BST 被修改的很頻繁，而且查找第k小元素的操作也很頻繁，問我們如何優化。其實最好的方法還是像上面的解法那樣利用分治法來快速定位目標所在的位置，但是每個遞歸都遍歷左子樹所有結點來計算個數的操作並不高效，所以應該修改原樹結點的結構，使其保存包括當前結點和其左右子樹所有結點的個數，這樣就可以快速得到任何左子樹結點總數來快速定位目標值了。定義了新結點結構體，然后就要生成新樹，還是用遞歸的方法生成新樹，注意生成的結點的 count 值要累加其左右子結點的 count 值。然后在求第k小元素的函數中，先生成新的樹，然后調用遞歸函數。在遞歸函數中，不能直接訪問左子結點的 count 值，因為左子節結點不一定存在，所以要先判斷，如果左子結點存在的話，那么跟上面解法的操作相同。如果不存在的話，當此時k為1的時候，直接返回當前結點值，否則就對右子結點調用遞歸函數，k自減1，參見代碼如下：
解法四：
// Follow up class Solution { public: struct MyTreeNode { int val; int count; MyTreeNode *left; MyTreeNode *right; MyTreeNode(int x) : val(x), count(1), left(NULL), right(NULL) {} }; MyTreeNode* build(TreeNode* root) { if (!root) return NULL; MyTreeNode *node = new MyTreeNode(root->val); node->left = build(root->left); node->right = build(root->right); if (node->left) node->count += node->left->count; if (node->right) node->count += node->right->count; return node; } int kthSmallest(TreeNode* root, int k) { MyTreeNode *node = build(root); return helper(node, k); } int helper(MyTreeNode* node, int k) { if (node->left) { int cnt = node->left->count; if (k <= cnt) { return helper(node->left, k); } else if (k > cnt + 1) { return helper(node->right, k - 1 - cnt); } return node->val; } else { if (k == 1) return node->val; return helper(node->right, k - 1); } } };
Github 同步地址：
https://github.com/grandyang/leetcode/issues/230
類似題目：
Second Minimum Node In a Binary Tree
參考資料：
https://leetcode.com/problems/kth-smallest-element-in-a-bst/