文章22
标签43
分类12

TopK 问题两种解决方案:堆、快排变形

这里的 TopK 不单指最大的 kk 个元素,也可以是最小的 kk 个元素,具体的例题可以看 剑指 Offer 40. 最小的k个数。本文的代码均为可以在力扣提交的题解,故不包含头文件和 main 函数。

看到这种题,最先想到的可能是直接使用已有的排序函数,但是其时间复杂度为 O(nlogn)O(n \log n),大量的排序工作是多余。

一个优化的想法是,建立一个大小为 kk 的堆(针对求最小的 kk 个元素使用最大堆),只需要一次遍历将遍历到的元素和堆顶元素比较,若比堆顶小,则替换堆顶,然后进行堆调整。遍历完成后,堆中就是要求的最小的 kk 个元素(无序的)。

通过堆优化后时间复杂度为 O(nlogk)O(n \log k)

class Solution {
public:
    vector<int> getLeastNumbers(const vector<int> &arr, int k) {
        vector<int> ans(k, 0);
        if (k == 0) {
            return ans;
        }
        // 建堆
        for (int i = 0; i < k; ++i) {
            ans[i] = arr[i];
        }
        for (int i = k / 2 - 1; i >= 0; --i) {
            adjust(ans, i);
        }
        int n = arr.size();
        for (int i = k; i < n; ++i) {
            // 比堆中最大的要小
            if (arr[i] < ans[0]) {
                // 替换
                ans[0] = arr[i];
                // 调整后堆顶是最大的元素
                adjust(ans, 0u);
            }
        }
        return ans;
    }


    void adjust(vector<int> &heap, uint32_t index) {
        uint32_t t = index;
        // 左孩子
        if ((index << 1u) + 1 < heap.size() && heap[(index << 1u) + 1] > heap[t]) {
            t = (index << 1u) + 1;
        }
        // 右孩子
        if ((index << 1u) + 2 < heap.size() && heap[(index << 1u) + 2] > heap[t]) {
            t = (index << 1u) + 2;
        }
        // 孩子节点比父节点大
        if (t != index) {
            std::swap(heap[index], heap[t]);
            // 递归调整
            adjust(heap, t);
        }
    }
};

快排

这里不是真的快排,而是借助快排的思想,将元素划分成比中间元素后大或小的两部分后,再对其中的一个部分再次划分。快排的思想是分治,这里解决 TopK 的思想是减治,因为划分后不符合条件的部分直接舍弃掉。

该方法的期望时间复杂度 O(n)O(n),最坏时间复杂度 O(n2)O(n^2),因为可能每次选取的中间值都是最大值,为了解决这个问题,可以随机选择中间值(被注释掉的三行代码)。

class Solution {
public:
    std::vector<int> getLeastNumbers(std::vector<int> &arr, int k) {
        uint32_t left = 0, right = arr.size() - 1;
//        srand(time(nullptr));
        while (left < right) {
//            int t = rand() % (right - left + 1) + left;
//            std::swap(arr[left], arr[t]);
            int pivot = arr[left];
            int i = left, j = right;
            // 完成后以 i 为分界,左半部分比 arr[i] 小,右半部分比 arr[i] 大
            while (i <= j) {
                // 寻找左半部分比 pivot 大的元素
                while (i <= j && arr[i] < pivot) {
                    ++i;
                }
                // 寻找右半部分比 pivot 小的元素
                while (i <= j && arr[j] > pivot) {
                    --j;
                }
                if (i <= j) {
                    std::swap(arr[i], arr[j]);
                    ++i;
                    --j;
                }
            }
            if (i >= k) {
                // 大于等于 k 个
                right = i - 1;
            } else {
                // 小于 k 个
                left = i;
            }
        }
        return std::vector<int>(arr.begin(), arr.begin() + k);
    }
};

两种算法的选择

在数据量比较小时,两个算法没有太大的区别。

但是当数据量达到十亿级别(腾讯面试被问过这个问题)使用快排需要将所有数据读入内存(当然也可以优化,但是会更复杂)。

而使用堆只需要维护一个大小为 kk 的堆(实际来说是数组),由于只需要遍历一次数据,数据读入一次就行了,对磁盘的操作可以很少。相比 CPU 的速度,磁盘 I/O 的时间有数量级上的差异。

因此在大数据量上应该选择堆。

本文作者:Dreace
本文链接:https://blog.dreace.top/2020/Two-Solutions-to-the-Topk-Problem-Heap-Quick-Sort/
版权声明:本文采用 CC BY-NC-SA 3.0 CN 协议进行许可