LeetCode第 703 题：数据流中的第K大元素(C++)

703. 数据流中的第K大元素 - 力扣（LeetCode）

其实就是堆的思路，和这个很像：

LeetCode第 215 题：数组中的第K个最大元素(C++)_qq_32523711的博客-CSDN博客

不过虽然是简单题，不过题目读起来可一点都不简单，把数据流其实就是对应数组nums，比如例子中的[4, 5, 8, 2]，add函数会往流（数组）里面添加数字，添加必然造成流里面元素顺序发生变化，所以流里面的元素是动态更新的，而每次都需要找到流里面的第k大的数字。

首先肯定会想到排序，但是由于数据动态更新，每次添加一个元素进来都进行排序的话，时间复杂度会很高。

那么就像上面那个链接说的，查找第K大元素，主要有两种思路，类似快速排序的划分思路和堆（优先级队列）的思路。

由于这一题里面数据是动态更新的，如果使用划分的话，每次返回第k大元素的复杂度都是O(n)，n为数组元素个数（而且n会慢慢增大）。

堆的思路：

我们可以维护一个大小为K的小顶堆（堆顶元素为堆中最小值），当新进的元素比堆顶元素小时，就什么也不做，当比堆顶元素大时，就删除堆顶元素，将该元素入堆。当数组遍历完成之后，堆顶元素就是我们需要的第K个最大的元素。

class KthLargest {
public:
    KthLargest(int k, vector<int>& nums) {
        for(const auto &c : nums){
            q.push(c);
            if(q.size() > k) q.pop();
        }
        index = k;
    }
    
    int add(int val) {
        q.push(val);
        if(q.size() > index)    q.pop();
        return q.top();
    }
private:
    int index;
    priority_queue<int, vector<int>, greater<int>> q; //小顶堆
};

上述思路可以优化如下：如果新添加的数比堆顶的数要小的话，是不会改变第k大的数是哪个的，所以就不需要往堆里面添加，如果比堆顶大的话，就需要添加。不过需要注意堆是否已满：

class KthLargest {
public:
    KthLargest(int k, vector<int>& nums) {//nums可能为空
        for(int i = 0; i < k && i < nums.size(); ++i)  q.push(nums[i]);
        for(int i = k; i < nums.size(); ++i){
            if(nums[i] <= q.top()) continue;
            q.pop();
            q.push(nums[i]);
        }
        size = k;
    }

    int add(int val) {
        if(q.size() < size) q.push(val);
        else{
            if(val <= q.top()) return q.top();
            q.pop();
            q.push(val);
        }
        return q.top();
    }
private:
    int size;
    priority_queue<int, vector<int>, greater<int>> q; //小顶堆
};

如果需要自己去实现小顶堆，而不用优先级队列的话，这样就能过，但是，我这个代码其实是不完善的，严格来说是错的，但是提交能够通过。说明这个题目的测试用例实在不够完善，我这个代码错误的地方在于add函数的第一个if里面，在添加了一个元素之后（该元素会添加在数组尾部），但是我想当然地对数组首元素进行下沉堆化（自上而下），其实这一步是错的，随便找个例子：

["KthLargest","add","add","add","add","add"]
[[5,[4,5,8,2]],[1],[5],[10],[9],[4]]

if语句在第一次add的时候会加入元素“1”,那么此时数组里的元素应该是1,2,4,5,8，第k（5）大的元素显然是“1”,但是我的代码是对首元素下沉最后也是返回首元素，答案是“2”，显然是代码有误，实际上我应该做的是将加入在数组末尾的元素进行自下而上的堆化。

所以这个题目的测试用例还是不完善的。

class KthLargest {
public:
    void buildHeap(vector<int> &a, int n){
        for(int i = (n-1)/2; i >= 0; --i)   heapify(a, n, i);
    }
    void heapify(vector<int> &a, int n, int i){
        while(true){
            int maxPos = i;
            if(2*i+1 <= n && a[i] > a[2*i+1])   maxPos = 2*i+1;
            if(2*i+2 <= n && a[maxPos] > a[2*i+2])   maxPos = 2*i+2;
            if(i == maxPos) break;
            swap(a[i], a[maxPos]);
            i = maxPos;
        }
    }
    KthLargest(int k, vector<int>& nums) {
        for(int i = 0; i < k && i < nums.size(); ++i)  a.push_back(nums[i]);
        buildHeap(a, a.size()-1); //先导入k个元素并建堆
        for(int i = k; i < nums.size(); ++i){//剩下的元素插入需要进行堆化
            if(nums[i] < a[0])  continue;
            a[0] = nums[i]; //直接复制给堆顶元素（堆顶就是最小的）
            heapify(a, a.size()-1, 0);//堆顶元素下沉（堆化）
        }
        size = k;
    }
    int add(int val) {
        if(a.size() < size){ //nums 的长度≥ k-1，所以这一步最多就执行一次
            a.push_back(val);
            heapify(a, a.size()-1, 0);
        } 
        else{
            if(val <= a[0]) return a[0];
            a[0] = val;
            heapify(a, size-1, 0);
        }
        return a[0];
    }
private:
    int size;
    vector<int> a;
};

下面才是正确的代码，定义insert函数插入，然后自下而上堆化：

class KthLargest {
public:
    void buildHeap(vector<int> &a, int n){
        for(int i = (n-1)/2; i >= 0; --i)   heapify(a, n, i);
    }
    void heapify(vector<int> &a, int n, int i){
        while(true){
            int maxPos = i;
            if(2*i+1 <= n && a[i] > a[2*i+1])   maxPos = 2*i+1;
            if(2*i+2 <= n && a[maxPos] > a[2*i+2])   maxPos = 2*i+2;
            if(i == maxPos) break;
            swap(a[i], a[maxPos]);
            i = maxPos;
        }
    }
    void insert(int val){
        a.push_back(val);
        int i = a.size()-1; //尾元素下标
        //自下而上堆化
        while( a[i] < a[(i-1)/2]){//如果该元素比它的父节点要小，那它就应该“浮上去”
            swap(a[i], a[(i-1)/2]);
            i = (i-1)/2;
        }
    }
    KthLargest(int k, vector<int>& nums) {
        for(int i = 0; i < k && i < nums.size(); ++i)  a.push_back(nums[i]);
        buildHeap(a, a.size()-1); //先导入k个元素并建堆
        for(int i = k; i < nums.size(); ++i){//剩下的元素插入需要进行堆化
            if(nums[i] < a[0])  continue;
            a[0] = nums[i]; //直接复制给堆顶元素（堆顶就是最小的）
            heapify(a, a.size()-1, 0);//堆顶元素下沉（堆化）
        }
        size = k;
    }
    int add(int val) {
        if(a.size() < size){ //nums 的长度≥ k-1，所以这一步最多就执行一次
            //a.push_back(val);
            //heapify(a, a.size()-1, 0);
            insert(val);
        } 
        else{
            if(val <= a[0]) return a[0];
            a[0] = val;
            heapify(a, size-1, 0);
        }
        return a[0];
    }
private:
    int size;
    vector<int> a;
};

用其他的数据结构其实也是可以的，关键在于处理好排序，可以使用multiset，自身可以排序：

class KthLargest {
public:
    KthLargest(int k, vector<int>& nums) {
        for(const auto &c : nums){
            set.insert(c);
            if(set.size() > k) set.erase(set.begin());;
        }
        index = k;
    }
    
    int add(int val) {
        set.insert(val);
        if(set.size() > index)    set.erase(set.begin());
        return *set.begin();
    }
private:
    int index;
    multiset<int> set;
};