跳转至

703.数据流中的第 K 大元素 (Easy)*

题目描述*

设计一个找到数据流中第K大元素的类（class）。注意是排序后的第K大元素，不是第K个不同的元素。

你的 KthLargest 类需要一个同时接收整数 k 和整数数组nums 的构造器，它包含数据流中的初始元素。每次调用 KthLargest.add，返回当前数据流中第K大的元素。

示例*

int k = 3;
int[] arr = [4,5,8,2];
KthLargest kthLargest = new KthLargest(3, arr);
kthLargest.add(3);   // returns 4
kthLargest.add(5);   // returns 5
kthLargest.add(10);  // returns 5
kthLargest.add(9);   // returns 8
kthLargest.add(4);   // returns 8

说明*

你可以假设 nums 的长度≥ k-1 且k ≥ 1。

代码*

看着应该是用堆实现比较简单，不过这个题是在探索二叉搜索树里的，所以先用 BST 实现以下。还可以使用 multiset 自动排序。应该还是堆最快。

class KthLargest {
private:
    TreeNode *root = nullptr;
    unordered_map<TreeNode*, int> nodeCnt;
    int kth;
    TreeNode* insert(TreeNode *root, int _val) {
        if(root == nullptr) {
            root = new TreeNode(_val);
            nodeCnt[root] = 1;
            return root;
        }
        if(_val < root->val) {
            root->left = insert(root->left, _val);
        }else if(_val > root->val) {
            root->right = insert(root->right, _val);
        }
        nodeCnt[root] += 1;
        return root;
    }
    TreeNode* searchKth(TreeNode *root, int k) {
        if(root == nullptr) {
            return root;
        }
        int leftCnt = (root->left == nullptr ? 0 : nodeCnt[root->left]);
        int rightCnt = (root->right == nullptr ? 0 : nodeCnt[root->right]);
        int rootCnt = nodeCnt[root] - leftCnt - rightCnt;
        if(k <= rightCnt) {
            return searchKth(root->right, k);
        }else if(k > rightCnt + rootCnt) {
            return searchKth(root->left, k - rightCnt - rootCnt);
        }else {
            return root;
        }
    }
public:
    KthLargest(int k, vector<int>& nums) {
        kth = k;
        for(auto i : nums) {
            root = insert(root, i);
        }
    }

    int add(int val) {
        root = insert(root, val);
        TreeNode *res = searchKth(root, kth);
        return res->val;
    }
};

class KthLargest {
private:
    int k;
    multiset<int> s;
public:
    KthLargest(int _k, vector<int>& nums) {
        for(auto i : nums) {
            s.insert(i);
            if(s.size() > _k) {
                s.erase(s.begin());
            }
        }
        k = _k;
    }

    int add(int val) {
        s.insert(val);
        if(s.size() > k) {
            s.erase(s.begin());
        }
        return *s.begin();
    }
};

class KthLargest {
private:
    priority_queue<int, vector<int>, greater<int>> q;
    int k;
public:
    KthLargest(int k, vector<int>& nums) {
        for(auto i : nums) {
            q.push(i);
            if(q.size() > k) {
                q.pop();
            }
        }
        this->k = k;
    }

    int add(int val) {
        q.push(val);
        if(q.size() > k) {
            q.pop();
        }
        return q.top();
    }
};
最后更新: July 23, 2022