TopK问题详解

题目描述

面试中经常会问到的一道题目：从n个未排序的数中得到的最大的k个数，称为TopK问题。（最小的k个数做法也相似）

基于partition函数

基于快速排序中的partition函数，时间复杂度为O（n），空间复杂度为O（1）；需要改变输入；

（1）根据Partition函数得到索引值index，index前的数据均大于nums[index]，index后的数据均小于nums[index]

（2）如果index = k-1，则已经划分完成；数组前k个数据即为最大的k个数

（3）如果index>k-1，begin=index+1；否则end = index-1

（4）重复（2）操作直到index = k-1

class SolutionI {
public:
    /*
     * 基于快排Partition函数
     * 时间复杂度O（n）空间复杂度O（1）需要改变输入
     */
    vector<int> getTopK(vector<int> nums, int k){
        if (nums.empty() || k > nums.size() || k<=0) return {};
        vector<int> ret;
        int begin = 0, end = nums.size()-1;
        int idx = Partition(nums,begin,end);
        while (idx != k-1){
            if (idx < k-1){
                begin = idx + 1;
                idx = Partition(nums,begin,end);
            }else{
                end = idx - 1;
                idx = Partition(nums,begin,end);
            }
        }

        for (int i = 0; i < k; ++i) {
            ret.push_back(nums[i]);
        }
        return ret;
    }

    // 返回索引值idx，idx前的元素均大于该处的元素值；idx后的元素均小于该处的元素值
    int Partition(vector<int> &nums, int begin, int end){
        if (begin > end) return begin;
        int key = nums[begin];    // 取最后一个值为基准值
        while (begin < end){
            while (nums[end] <= key && begin < end) --end;
            nums[begin] = nums[end];
            while (nums[begin] > key && begin < end) ++ begin;
            nums[end] = nums[begin];
        }
        nums[begin] = key;
        return begin;
    }
};

堆

动态维护大小为k的堆，时间复杂度为O（nlogk），空间复杂度为O（k），无需改变输入；

用multiset实现最小堆；其基于红黑树实现，查找、删除、插入时间复杂度为O（logk）

每次比较最小堆的堆顶元素（topK大元素中最小值）与当前元素，如果当前元素更大，替换该元素。

class SolutionII {
public:
    /*
     * 最小堆方法
     * 时间复杂度O（nlogk）空间复杂度O（k）
     */
    vector<int> getTopK(vector<int>& nums, int k){
        if (nums.empty() || k > nums.size() || k<1) return {};
        vector<int> ret;
        multiset<int, less<int>> m; // 最小堆，保存最大的K个数
        multiset<int, less<int>>::iterator it;
        for (int i = 0; i < nums.size(); ++i) {
            if(m.size()<k){
                m.insert(nums[i]);
            } else{
                it = m.begin();
                // 如果当前值大于topK的最小元素（最小堆堆顶），替换该值
                if (nums[i]> *it){
                    m.erase(it);
                    m.insert(nums[i]);
                }
            }
        }
        for (it = m.begin();it!=m.end();it++)
            ret.push_back(*it);

        return ret;
    }
};

两种算法的比较

	基于Partition的解法	基于堆或红黑树
时间复杂度	O（n）	O（nlogk）
是否需要修改输入数据	是	否
是否适用于海量数据	否	是

基于partition函数的平均时间复杂度更快，但是具有明显的限制，如会修改输入数组

基于堆或红黑树的方法，具有两个明显优点：（1）没有修改输入数据（2）适合海量数据的输入。由于内存大小有限，可能无法一次性将所有数据全部载入内存，这时可每次读入一个数字，再进行判断；只需要内存能够容纳k个数据即可，适用于n很大并且k很小的问题。