【TopK】O(n)复杂度Topk问题的三种解决方法（优先级队列、排序、快速选择）

一、题目

Top-K问题是从一组数据中找出排名前K的元素。这类问题在搜索引擎、推荐系统、数据分析及数据挖掘等多个领域有着广泛应用。该问题通常需要高效地处理大规模数据，减少计算时间和内存需求。

在解决Top-K问题时，一种有效的方法是利用堆这种数据结构。通过构建合适大小的堆（大根堆或小根堆），可以有效地找到前K个最大或最小的元素。具体步骤包括先用数据集合中的前K个元素建堆，然后将剩余的元素与堆顶元素比较，必要时进行替换，并维护堆的性质。

最小的K个数_牛客题霸_牛客网

寻找第K大_牛客题霸_牛客网

二、题解

优先级队列

优先级队列在解决topk问题时，如果要求最小的前k个或者最小的第k个，我们需要建立一个大根堆，首先将前k个元素入堆，然后从第k个开始进行判断，如果当前元素比堆顶的元素小就让他入堆，这样保证堆中每次都存这最小的k个里面最大的那一个，也就是前第k大的。如果是求最大的k个或者第k大的就需要创建一个小根堆，如果当前元素比堆顶元素大就弹出堆顶并让该元素入堆

#include<iostream>
using namespace std;
const int N = 100010;
int n, k, q[N], heap[N], sz;

void siftDown(int parent)
{
    int child = (parent << 1) + 1;
    while (child < sz)
    {
        if (child + 1 < sz && heap[child + 1] > heap[child]) child++;
        if (heap[child] > heap[parent]) 
        {
            int t = heap[child]; heap[child] = heap[parent]; heap[parent] = t;
            
            parent = child;
            child = (parent << 1) + 1;
        } else break;
    }
}

void siftUp(int child)
{
    int parent = (child - 1) >> 1;
    while (parent >= 0)
    {
        if (heap[child] > heap[parent])
        {
            int t = heap[child]; heap[child] = heap[parent]; heap[parent] = t;
            
            child = parent;
            parent = (child - 1) >> 1;
        } else break;
    }
}

void heapify() 
{
    for (int i = (sz - 1 - 1) >> 1; i >= 0; i--) siftDown(i);
}

void pop()
{
    int t = heap[0]; heap[0] = heap[sz - 1]; heap[sz - 1] = t;
    sz--;
    siftDown(0);
}

void push(int x)
{
    if (sz > k) return;
    heap[sz] = x;
    siftUp(sz++);
}

void topK()
{
    // 1. 建立大根堆
    heapify();
    // 2. 入堆
    for (int i = 0; i < k; i++) push(q[i]);
    for (int i = k; i < n; i++)
        if (q[i] < heap[0]) {
            pop(); push(q[i]);
        }
}

int main()
{
    cin >> n >> k;
    for (int i = 0; i < n; i++) cin >> q[i];
    topK();
    cout << heap[0] << endl;
    return 0;
}

排序

利用快排的思想排序后获取第k小或者第k大的元素

#include<iostream>
using namespace std;
const int N = 100010;
int n, k, q[N];

void qsort(int l, int r) 
{
    if (l >= r) return;
    srand(time(0));
    int i = l, s = l - 1, e = r + 1, x = q[rand() % (r - l + 1) + l];
    while (i < e) 
    {
        if (q[i] > x) swap(q[i], q[--e]);
        else if (q[i] < x) swap(q[i++], q[++s]);
        else i++;
    }
    qsort(l, s), qsort(e, r);
}

int main()
{
    cin >> n >> k;
    for (int i = 0; i < n; i++) cin >> q[i];
    qsort(0, n - 1);
    cout << q[k - 1] << endl;
    return 0;
}

快速选择

上述快排代码每一次递归都将数组的l-r区间分为3块，l - s 元素都会小于x， s + 1 - e-1 中的元素都是等于x的，e - r中的元素都是大于x的，根据这个性质我们在每次递归的时候就可以通过比较去控制下一场要去哪一个区间去寻找第k小的元素，根据上述思想，将 三个区间元素个数分别标记为a，b，c，a = s - l + 1, c = r - e + 1, b = e- s - 1，a就是在这个区间中最小的前a个元素，c就是在这个区间里面最大的b个元素，如果在一次递归中当a的值大于等于k的时候说明第k小的元素在前a个元素当中那么下一次就到l - s 区间中去寻找，如果a小于k但是a+b是大于等于k的，说明第k小的元素在中间这个区域里，而中间这些元素都是x，也就是说第k小的元素就是x，如果上述都不满足，说明第k小的元素在右边这段区域内，我们只需要在右边区域中去寻找k - a -b小的元素

#include<iostream>
using namespace std;
const int N = 100010;
int n, k, q[N];

int qsort(int l, int r, int k) 
{
    if (l > r) return -1;
    srand(time(0));
    int i = l, s = l - 1, e = r + 1, x = q[rand() % (r - l + 1) + l];
    while (i < e) 
    {
        if (q[i] > x) swap(q[i], q[--e]);
        else if (q[i] < x) swap(q[i++], q[++s]);
        else i++;
    }
    int a = s - l + 1, b = e - s - 1;
    if (a >= k) return qsort(l, s, k);
    else if (a + b >= k) return x;
    else return qsort(e, r, k - a - b);
}

int main()
{
    cin >> n >> k;
    for (int i = 0; i < n; i++) cin >> q[i];

    cout << qsort(0, n - 1, k) << endl;
    return 0;
}