二叉搜索树中第K小的元素
我的解答:
public int kthSmallest(TreeNode root, int k) {
Deque<TreeNode> stack = new LinkedList<>();
int cnt = 0;
while(!stack.isEmpty() || root != null){
while(root != null){
stack.push(root);
root = root.left;
}
root = stack.poll();
cnt++;
if(cnt == k){
return root.val;
}
root = root.right;
}
return 0;
}
分析:
1、代码的时间复杂度为O(height+k),空间复杂度为O(height),height为树的高度。解题思路:利用"搜索二叉树的中序遍历是升序序列"的性质,利用栈对搜索二叉树进行升序遍历,并记录中序遍历到的节点个数,当遍历到第k个节点时,这个节点的值就是答案。
2、进阶:如果二叉搜索树经常被修改(插入/删除操作)并且你需要频繁地查找第 k 小的值,你将如何优化算法?
答:初步想法是维护一个双向链表,该链表记录了搜索二叉树中序遍历的顺序,每次插入或删除元素,我们直接根据插入或删除的节点改变链表节点之间的联系即可。(看了官方题解后,知道需要自己实现平衡二叉搜索树(AVL树)来实现)
看了官方题解后的解答:
//方法一:中序遍历
//时间复杂度:O(H+k),H为树的高度,在开始遍历之前,我们需要O(H)到达叶结点
//空间复杂度:O(H),H为树的高度
public int kthSmallest(TreeNode root, int k) {
Deque<TreeNode> stack = new LinkedList<>();
while(!stack.isEmpty() || root != null){
while(root != null){
stack.push(root);
root = root.left;
}
root = stack.poll();
k--;
if(k == 0){
break;
}
root = root.right;
}
return root.val;
}
//方法二(只看思路自己实现版,不太正确):记录子树的节点数(解决问题:如果你需要频繁地查找第 k 小的值,你将如何优化算法?)
//时间复杂度:O(H+k),H为树的高度,在开始遍历之前,我们需要O(H)到达叶结点
//空间复杂度:O(H),H为树的高度
public int kthSmallest(TreeNode root, int k) {
//key:当前节点 value:当前节点的左子树的节点个数
Map<TreeNode,Integer> map = new HashMap<>();
Deque<TreeNode> stack = new LinkedList<>();
int cnt = 0;//当前节点遍历之前,中序遍历到的节点个数,即当前节点左子树的节点个数
TreeNode cur = root;
while(!stack.isEmpty() || cur !=null){
while(cur!=null){
stack.push(cur);
cur = cur.left;
}
cur = stack.poll();
map.put(cur,cnt);
cnt++;
cur = cur.right;
}
cur = root;
while(cur != null){
if(map.get(cur)< k-1){
cur = cur.right;
}
else if(map.get(cur) > k-1){
cur = cur.left;
}
else{
break;
}
}
return cur.val;
}
//方法二(看了官方实现版):记录子树的节点数(解决问题:如果你需要频繁地查找第 k 小的值,你将如何优化算法?)
//时间复杂度:O(H+k),H为树的高度,在开始遍历之前,我们需要O(H)到达叶结点
//空间复杂度:O(H),H为树的高度
class Solution {
public int kthSmallest(TreeNode root, int k) {
MyBst myBst = new MyBst(root);
return myBst.kthSmallest(k);
}
}
class MyBst{
private TreeNode root;
private Map<TreeNode,Integer> nodeNum;
public MyBst(TreeNode root){
this.root = root;
this.nodeNum = new HashMap<>();
countNodeNum(root);
}
//计算以每个节点作为根节点的子树节点总数
private Integer countNodeNum(TreeNode cur){
if(cur == null){
return 0;
}
nodeNum.put(cur, 1 + countNodeNum(cur.left) + countNodeNum(cur.right));
return nodeNum.get(cur);
}
//获取以当前节点cur为根节点的子树节点总数
private Integer getNodeNum(TreeNode cur){
return nodeNum.getOrDefault(cur,0);
}
//获取搜索二叉树第k小的元素
public Integer kthSmallest(int k){
TreeNode cur = root;
int leftNum;
while(cur != null){
leftNum = getNodeNum(cur.left);
if(leftNum < k-1){
cur = cur.right;
k = k-leftNum-1;
}
else if(leftNum > k-1){
cur = cur.left;
}
else{
break;
}
}
return cur.val;
}
}
//方法三:平衡二叉搜索树(解决问题:如果二叉搜索树经常被修改(插入/删除操作)并且你需要频繁地查找第k小的值,你将如何优化算法?)
//时间复杂度:预处理的时间复杂度为 O(N),其中 N 是树中结点的总数。插入、删除和搜索的时间复杂度均为 O(logN)
//空间复杂度:O(N),用于存储平衡二叉搜索树
分析:
1、官方题解方法一的解题思路与我的解答思路一致,只是具体实现略有不同。
2、方法二的解题思路:设计一个新的类MyBst,计算并用哈希表记录二叉树以每个节点作为根节点时的子树节点总数,并实现获取搜索二叉树第k小的元素方法kthSmallest。用root构建MyBst类的对象,直接调用MyBst类对象的kthSmallest方法即可。
在知道二叉树以每个节点作为根节点时的子树节点总数后,如何快速获取搜索二叉树第k小的元素?我们知道每个节点的左子树节点总数leftNum,若leftNum>k-1,那么第k小的元素一定在当前节点的右子树上,我们继续往当前节点的右子树遍历,且k变为k-leftNum-1;若leftNum<k-1,那么第k小的元素一定在当前节点的左子树上,我们继续往当前节点的左子树遍历;若leftNum==k-1,那么当前节点就是第k小的元素。
3、官方题解三需要自己实现平衡二叉搜索树(AVL树),用于解决"二叉搜索树经常被修改(插入/删除操作)并且需要频繁地查找第k小的值"的问题,编码难度很大,暂时略过。
总结
- 本题主要掌握"搜索二叉树的中序遍历是升序序列"这个性质,根据这个性质即可解题。
- 对于本题的进阶问题"如果二叉搜索树经常被修改(插入/删除操作)并且你需要频繁地查找第 k 小的值,你将如何优化算法?",暂时只学习了进阶问题的一部分"搜索二叉树不变,但需要频繁地查找第 k 小的值",可利用哈希表维护每个节点作为根节点时子树的节点总数来实现快速查找。