「前言」文章内容是关于 TopK 问题的讲解。

一、概念

TopK 问题是一类常见的算法和数据处理问题，其核心任务是从包含大量数据项的集合中找到前 K 个最大的元素或者前 K 个最小的元素。

这类问题在计算机科学和数据处理领域广泛存在，具有多种应用场景，如搜索引擎、推荐系统、各种排行榜等等。

解决 TopK 问题的算法多种多样，常见的包括堆排序、快速选择等。这里只介绍使用堆排序进行解决 TopK 的问题。

因为快速排序不适合 TopK 海量数据的情况，而堆排序则可以，所以堆排序是解决 TopK 问题的较优选择。

堆排序的时间复杂度又是 O(N*logN)，这里就不再介绍堆排序了。

二、解决

求前 K 个最大元素或者前 K 个最小元素，只需建立 K 个数的堆即可。

即先建一个 K 个数的堆，然后将数组中 N-K 个元素依次与堆顶的元素比较，若比堆顶元素大（小），则将堆顶元素删除，然后将这个数插入到堆中，直到遍历完剩下的数据。遍历结束后，堆里面便是最大的前 K 个元素或者最小的前 K 个元素。

==注意==：求前 K 个最大元素建小堆，求前 K 个最小元素建大堆。

注意区分：堆排序排升序是建大堆，排降序是建小堆。

为什么求前 K 个最大元素建小堆？？

  • 其实很简单，假设我们是建大堆，万一在建堆的时候最大的元素就已经在堆顶了。后么不论怎么比较，你后面的 N-K 个元素都没法进堆，如果 N-K 个元素里面有其他第二大、第三大...的元素，压根进不了堆里面。
  • 同理求前 K 个最小元素建大堆。

2.1 C语言代码

C语言代码如下：（求前 K 个最大元素建小堆）

void Swap(int* p1, int* p2)
{
    int tmp = *p1;
    *p1 = *p2;
    *p2 = tmp;
}

// 堆的向下调整（下面是小堆）
void AdjustDown(int* arr, int n, int parent) // n:arr数组的大小; parent:父节点的数组下标
{
    // 下标以0开始
    // 左子节点下标为parent * 2 + 1; 右子节点的下标为parent * 2 + 2
    int child = parent * 2 + 1; // 假设左孩子较大
    // 
    while (child < n)
    {
        // 选出左右孩子中小的那个
        if ((child + 1 < n) && arr[child] > arr[child + 1]) // child + 1:右孩子的下标；-- 右孩子存在，并且左孩子比右孩子大
        {
            ++child;
        }
        // 孩子跟父亲比较
        if (arr[child] < arr[parent])
        {
            Swap(&arr[child], &arr[parent]); // 交换数据
            //迭代
            parent = child;
            child = parent * 2 + 1;
        }
        else
        {
            break;
        }
    }
}

// 求前k个大的元素
int* FindTopK(int* arr, int n, int k) // n:数组大小
{
    if (arr == NULL || k == 0 || k > n) return NULL;

    // 1. 用数组的前K个数建小堆
    int* tmpArr = (int*)malloc(sizeof(int) * k);
    if (tmpArr == NULL) return NULL;
    for (int i = 0; i < k; ++i) tmpArr[i] = arr[i];

    // 2.建堆
    // 从第一个非叶子结点开始向下调整，一直到根
    for (int i = (n - 1 - 1) / 2; i >= 0; --i) // n-1:数组下标上限，(n-1-1)/2:孩子的父亲节点
    {
        AdjustDown(tmpArr, k, i);
    }

    // 3.调整（排序的过程）
    // 剩下的n-k个数依次与堆顶数据比较
    for (int i = k; i < n; ++i)
    {
        if (arr[i] > tmpArr[0])
        {
            tmpArr[0] = arr[i]; // 堆顶数据替换
        }
        // 对根进行一次向下调整
        AdjustDown(tmpArr, k, 0); 
    }
    return tmpArr;
}

运行结果如下：

image-20240921211223798

如果是求前 K 个最小元素则建大堆。修改一下判断条件即可。

// ...
if ((child + 1 < n) && arr[child] < arr[child + 1])
// ...
if (arr[child] > arr[parent])
// ...

2.2 C++代码

如果是C++，直接使用 STL 的优先级队列即可（priority_queue）。

假设是求前 K 个最小元素，建大堆。

vector<int> FindTopK(vector<int>& arr, int k)
{
    int n = arr.size();
    // priority_queue默认是大堆
    priority_queue<int> heap;
    // 只需保留前k个元素
    for (auto& x : arr)
    {
        heap.push(x);
        if (heap.size() > k) heap.pop();
    }
    // 提取数据
    vector<int> result;
    while (heap.size())
    {
        result.push_back(heap.top());
        heap.pop();
    }
    return result;
    return result;
}

如果是求前 K 个最大元素则建小堆。

vector<int> FindTopK(vector<int>& arr, int k)
{
    int n = arr.size();
    // priority_queue默认是大堆
    // priority_queue<int> head;
    // 建小堆
    priority_queue<int, std::vector<int>, std::greater<int>> heap;

    // 只需保留前k个元素
    for (auto& x : arr)
    {
        heap.push(x);
        if (heap.size() > k) heap.pop();
    }
    // 提取数据
    vector<int> result;
    while (heap.size())
    {
        result.push_back(heap.top());
        heap.pop();
    }
    return result;
}

--------------- END ---------------

「 作者 」 枫叶先生
「 更新 」 2024.9.21
「 声明 」 余之才疏学浅，故所撰文疏漏难免，
          或有谬误或不准确之处，敬请读者批评指正。