LeetCode 题解

积跬步，至千里。

最小的k个数

分类：剑指 Offer 标签： 堆 分治算法 发布于：2020-05-11 ID：offer_40

题目描述

来源于 https://leetcode-cn.com/

输入整数数组 arr ，找出其中最小的 k 个数。例如，输入4、5、1、6、2、7、3、8这8个数字，则最小的4个数字是1、2、3、4。

示例 1：

输入：arr = [3,2,1], k = 2
输出：[1,2] 或者 [2,1]

示例 2：

输入：arr = [0,1,2,1], k = 1
输出：[0]

限制：

  • 0 <= k <= arr.length <= 10000
  • 0 <= arr[i] <= 10000

解法一：使用 partition

快速排序中，使用切分的方法，将原数组分为两部分。前一部分小于 pivot，第二部分大于 pivot。如果切分点在下标 k-1 处，那么前 k 个元素不就是最小的 k 个元素了吗。

class Solution {
public:
    vector<int> getLeastNumbers(vector<int>& arr, int k) {
        if(k == 0 || arr.size() == 0) return {};

        int lo = 0, hi = arr.size();
        int i = partition(arr, lo, hi);
        while(i != k-1){
            if(i > k-1){
                hi = i;
            }else{
                lo = i+1;
            }
            i = partition(arr, lo, hi);
        }
        return vector<int>(arr.begin(), arr.begin()+k);
    }

    int partition(vector<int>& nums, int lo, int hi){
        int v = nums[lo];
        int i = lo, j = hi;
        while(true){
            while(++i < j && nums[i] < v);
            while(nums[--j] > v);
            if(i >= j){
                break;
            }
            std::swap(nums[i], nums[j]);
        }
        std::swap(nums[lo], nums[j]);
        return j;
    }
};

本题中有几个坑要注意避免，while(i != k-1) 能不能改成 while(i != k) 呢，其实无论 pivot 在 k 或 k+1 处，前 k 个元素都是最小的。但是万一 k == arr.size() 怎么办呢，此时 while(i != k) 会出问题，因此必须写成 while(i != k-1)。此时如果 k = 0，就会出问题，因此在一开始就排除掉 k = 0 的情况。

C++ 中有 nth_element 算法，这个算法的实现思路和上面的思路一致。

class Solution {
public:
    vector<int> getLeastNumbers(vector<int>& nums, int k) {
        nth_element(nums.begin(), nums.begin()+k, nums.end());
        return vector<int>(nums.begin(), nums.begin()+k);
    }
};

使用优先队列

使用优先队列记录下最小的 k 个元素。使用一个最大堆，当堆的大小为 k+1 的时候，就删除堆中最大的元素，如此，堆中就保留了最小的 k 个数。

class Solution {
public:
    vector<int> getLeastNumbers(vector<int>& nums, int k) {
        priority_queue<int> pq;

        for(int n: nums){
            pq.push(n);
            if(pq.size() == k + 1){
                pq.pop();
            }
        }

        vector<int> ret;
        while(!pq.empty()){
            ret.push_back(pq.top());
            pq.pop();
        }

        return ret;
    }
};