Top k 问题的两种经典解法

目录

方法一：堆

方法二：基于快排partition函数

2种方法比较

【题目】给定一个长为n的数组，找出其中最小的k个数（或 第k小的数）。

【思路】以找"最小的k个数"为例。

方法一：堆

class Solution {
public:
    vector<int> getLeastKNums(vector<int>& arr, int k) {
        int n = arr.size();
        vector<int> v;
        if (k<=0 || k>n) return v;
        priority_queue<int, vector<int>, less<int>> q;    // 大顶堆,也可直接写priority_queue<int>(默认less)

        for(int i=0; i<k; i++) 
            q.push(arr[i]);
        for(int i=k; i<n; i++){
            if (arr[i] < q.top()) {
                q.pop();
                q.push(arr[i]);
            }
        }

        for (int i = 0; i < k; i++) {
            v.push_back(q.top());
            q.pop();
        }
        return v;
    }
};

方法二：基于快排partition函数

回顾一下快速排序的思路。快速排序中有一步很重要的操作是 partition（划分），即从数组中随机选取一个数v作为基准数，然后原地移动数组中的元素，使得比 v 小的元素在 v 的左边，比 v 大的元素在 v 的右边。

我们的目的是寻找最小的 k 个数。假设经过一次 partition 操作后，基准数的下标为m，也就是说，左侧m+1个元素是原数组中最小的m+1个数。那么可以比较 m 和 k 的大小：

• 若m = k-1，我们就找到了最小的k个数，就是左侧[0,k-1]上的数；

• 若m > k-1，则最小的 k个数一定都在左侧数组中，我们只需要对左侧数组继续进行 partition；

• 若m < k-1，则左侧的 m 个数都属于最小的k个数，但还需要在右侧数组中寻找剩下的最小的k-m个数，所以对右侧数组继续进行partition 。

int partition(vector<int>& nums, int left, int right){  // 将数组分为两部分,并返回基准数的下标
    int pivot = nums[left];
    int i = left, j = right;
    while(i < j){
        // 这两个while只有一个被执行，因为另一个就是基准数pivot
        while((i<j) && (nums[j]>pivot)) // 不能用>=
            j--;
        while((i<j) && (nums[i]<pivot)) // 不能用<=
            i++;

        if((nums[i]==nums[j]) && (i<j))   // 避免两个数都等于pivot导致无限循环,如2,4,1,4,3,2
            i++;
        else swap(nums[i], nums[j]);
    }
    return i;   // i==j
}

vector<int> getLeastKNums(vector<int>& nums, int k){
    int n = nums.size();
    vector<int> v;
    if (k<=0 || k>n) return v;

    int start = 0, end = n-1;
    int index = partition(nums, start, end);
    while(index != k-1){
        if(index > k-1){    // 最小的k个数都在index左边
            end = index-1;
            index = partition(nums, start, end);
        }
        else{    // 在index右边寻找剩下的k-index个数
            start = index+1;
            index = partition(nums, start, end);
        }
    }

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

2种方法比较

参考：剑指Offer、https://mp.weixin.qq.com/s/je3GvIBPLGZz5Pt_TIQ3lA