C++的STL 堆 实现获取数组堆第K大的数

前言

堆数据结构 使用的是优先级队列实现，创建堆的时候需要指定堆中元素的排列方式，即最大堆或者最小堆
最大堆即 堆顶元素为堆中最大的元素
最小堆即 堆顶元素为堆中最小堆元素

如下为一个最大堆

---

回到文章标题，获取一个数组中第K大的数，要求时间复杂度是O(n)，即一次遍历即可获取到该值。
所以普通的先排序，再获取第k个数值的方式显然不满足要求，排序最快的快速排序也需要O(log2 n)

这里就是我们前言中提到的堆 数据结构堆优势了，可以构建一个最小堆，堆顶始终为堆中所有元素的最小值。当然在本次实现中，仅需要将堆堆大小限制为要求的K个即可，此时遍历完数组的元素，堆顶即为数组中第K大的值。

基本过程如下图：

实现如下(文末有测试代码)：

int calculte_k_maxnum(vector<int> &arr, int k) {
    priority_queue <int,vector<int>, greater<int>> small_heap;//有衔接队列构建最小堆

    if (arr.empty()) {
        print_erro("the array is empty\n",__LINE__);
    }
	
	/*按照步骤入堆*/
    for(int i =0;i < arr.size(); ++i) {
        if (small_heap.size() < k) {
            small_heap.push(arr[i]);
        } else if (small_heap.size() == k && small_heap.top() < arr[i]) {
            small_heap.pop();
            small_heap.push(arr[i]);
        }
    }
    /*最终的堆顶元素即为第K大的元素*/
    return small_heap.top();
}

测试代码如下：

#include 
#include 
#include 

using namespace std;

void print_erro(string s,int line) {
    cout << s <<" " << line << endl;
    exit(-1);
}

int calculte_k_maxnum(vector<int> &arr, int k) {
    priority_queue <int,vector<int>, greater<int>> small_heap;

    if (arr.empty()) {
        print_erro("the array is empty\n",__LINE__);
    }

    for(int i =0;i < arr.size(); ++i) {
        if (small_heap.size() < k) {
            small_heap.push(arr[i]);
        } else if (small_heap.size() == k && small_heap.top() < arr[i]) {
            small_heap.pop();
            small_heap.push(arr[i]);
        }
    }
    
    return small_heap.top();
}

int main(){
    vector<int> arr;
    int k;
    int tmp;

    cout << "input the number of arr you want to get max " << endl;
    cin >> k;

    cout << "input the arr number " << endl;
    for (int i = 0;i < 5; ++i) {
        cin >> tmp;
        arr.push_back(tmp);
    }

    cout << "the k's max number of arr is " << calculte_k_maxnum(arr,k) << endl;

    return 0;
}

输出如下:

input the number of arr you want to get max 
4
input the arr number 
2 4 3 6 7
the k's max number of arr is 3

input the number of arr you want to get max 
1        
input the arr number 
1 4 7 9 6
the k's max number of arr is 9