【剑指Offer-40】最小的k个数

问题

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

示例

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

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

解答1：快速排序思想

class Solution {
public:
    vector<int> getLeastNumbers(vector<int>& arr, int k) {
        srand(time(NULL));
        quicksort(arr, 0, arr.size() - 1, k);
        return vector<int>(arr.begin(), arr.begin() + k);
    }
private:
    void quicksort(vector<int>& nums, int left, int right, int k) {
        if (left >= right) return;
        int mid = partition(nums, left, right);
        // ------ 与快速排序不一样的地方
        if (mid > k - 1) quicksort(nums, left, mid - 1, k);
        else if (mid < k - 1) quicksort(nums, mid + 1, right, k);
        // ------
    }
    int partition(vector<int>& nums, int left, int right) {
        swap(nums[left], nums[rand() % (right - left + 1) + left]);
        int start = left, pivot = nums[left];
        while (left < right) {
            while (left < right && pivot <= nums[right]) right--;
            while (left < right && pivot >= nums[left]) left++;
            swap(nums[left], nums[right]);
        }
        swap(nums[start], nums[left]);
        return left;
    }
};

重点思路

快速排序过程中，能够保证pivot左边的一定小于pivot值，右边的一定大于。本题需要的是前k个数字，所以只需要当pivot的角标为k-1时，即可输出。所以我们在执行快速排序时，当pivot角标小于k-1时，递归右区间；大于k-1时，递归左区间，直到等于时输出。

解答2：最大堆思想

class Solution {
public:
    vector<int> getLeastNumbers(vector<int>& arr, int k) {
        priority_queue<int> p;
        vector<int> res;
        for (int i : arr) {
            if (p.size() < k || !p.empty() && p.top() > i) p.push(i);
            if (p.size() > k) p.pop();
        }
        while (!p.empty()) {
            res.push_back(p.top());
            p.pop();
        }
        return res;
    }
};

重点思路

C++的priority_queue默认为最大堆，即priority_queue<int, vector<int>, less<int>>，如果要使用最小堆的话则为priority_queue<int, vector<int>, greater<int>>。本题需要维护一个容量为k的最大堆。最大堆的堆顶为堆中元素的最大值。当堆中已经含有k个元素时，如果遍历得到的arr[i]小于堆顶元素，则替换堆顶元素。