题目描述
给定一个二叉搜索树,编写一个函数 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
函数?
解题思路
利用中序遍历的思想,首先获得左子树的元素总数,若加一正好等于k,则当前元素即为结果;若小于k则继续在右子树中寻找,注意让k减去左子树与根节点的元素之和。
代码
1 /**
2 * Definition for a binary tree node.
3 * struct TreeNode {
4 * int val;
5 * TreeNode *left;
6 * TreeNode *right;
7 * TreeNode(int x) : val(x), left(NULL), right(NULL) {}
8 * };
9 */
10 class Solution {
11 public:
12 int kthSmallest(TreeNode* root, int k) {
13 int res;
14 if(kth(root, k, res)) return res;
15 }
16 int kth(TreeNode* root, int k, int &res){
17 if(root == NULL) return 0;
18 int left = kth(root->left, k, res);
19 int right = 0;
20 if(left + 1 == k)
21 res = root->val;
22 else if(left < k)
23 right = kth(root->right, k - left - 1, res);
24 return left + right + 1;
25 }
26 };