Leetcode-230 Kth Smallest Element in a BST

发表于 | 分类于 | 次

Kth Smallest Element in a BST

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

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?

Example 1

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

解题思路

中序遍历，到达第 k 个节点就停止。

Follow up：中序遍历到第 k 个节点停止，必须先走到叶子节点，然后出栈至少 k 次，时间复杂度为 O(h+k)；插入删除的时间复杂度为O(h)；可以改进查找第 k 小的时间复杂度。借鉴 LRU 缓存中使用的哈希表和双向链表，可以在底层增加一个双向链表，BST 节点指向底层双向链表节点。插入时，根据插入位置的左或右在双向链表节点的前或后进行插入；删除时，直接从双向链表中删除节点并连接前后节点；查找第 k 小的时间复杂度为 O(k)。

复杂度分析

  • 时间复杂度：$O(h+k)$。
  • 空间复杂度：$O(h)$。

代码

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
class Solution {
public:
    int kthSmallest(TreeNode* root, int k) {
        stack<TreeNode*> sp;
        TreeNode *p = root;
        vector<int> res;
        while(!sp.empty() || p != nullptr){
            if(p) {
                sp.push(p);
                p = p->left;
            } else {
                res.push_back(sp.top()->val);
                if(res.size() == k) break;
                p = sp.top()->right;
                sp.pop();
            }
        }
        return res[k-1];
    }
};