leetcode：面试题 17.14. 最小K个数

题目来源

• leetcode：面试题 17.14. 最小K个数

题目描述

class Solution {
public:
    vector<int> smallestK(vector<int>& arr, int k) {

    }
};

题目解析

• 数据量： 10^5，所以O(N * logN)

大根堆

• 维护一个k长度的小根堆。 假如当前堆的大小已经有k，现在压入一个新元素，那么将会排除堆中最小的那个，所以小根堆中维护的是第1大的，第2大的，第3大的…
• 维护一个k长度的大根堆。假如当前堆的大小已经有k，现在压入一个新元素，那么将会排除堆中最大的那个，所以大根堆中维护的是第1小的，第2小的，第3小的…

因此，我们维护一个长度为k的大根堆。

• 遍历数组，并维护一个长度为k的大根堆
• 由提示知道，最终答案一定会满足，所以不用担心不足k个数。然后又可以任意顺序返回，所以最终我们直接将大根堆中的数据倒出来到vec中即可

class Solution {
public:
    vector<int> smallestK(vector<int>& arr, int k) {
        std::vector<int> ans;
        std::priority_queue<int, std::vector<int>, std::less<>> qMax;
        for (int & i : arr) {
            qMax.push(i);
            if(qMax.size() > k){
                qMax.pop();
            }
        }
        while (!qMax.empty()){
            ans.push_back(qMax.top());qMax.pop();
        }
        return ans;
    }
};

排序，然后取前k个数

class Solution {
public:
    vector<int> smallestK(vector<int>& arr, int k) {
        vector<int> vec(k, 0);
        sort(arr.begin(), arr.end());
        for (int i = 0; i < k; ++i) {
            vec[i] = arr[i];
        }
        return vec;
    }
};

快速排序

注意到，题目要求[任意顺序返回这k个数即可]，因此我们只需要保证前k小的数都出现在下标[0, k)的位置即可。

利用[快速排序]的数组划分即可做到。

class Solution {
    // 荷兰国旗，三色分区：
    // 在arr[l ... r]范围上，将小于pivot的数全部放到左边，等于pivot的全部放中间，大于pivot的全部放右边
    std::vector<int> partition(vector<int>& arr, int l, int r, int pivot){
        int less = l -1, more = r +1, curr = l;
        while (curr < more) {
            if (arr[curr] < pivot) {
                swap(arr[curr++], arr[++less]);
            } else if (arr[curr] == pivot) {
                curr++;
            } else {
                swap(arr[curr], arr[--more]);
            }
        }
        return {less + 1, more - 1};
    }
    // 在arr中，找到（如果排序的话）位于k-1位置（下标）的数：
    int findTheIndexNum(vector<int>& arr, int l,  int r, int idx){
        int  ans = 0;
        int pivot = arr[l];
        auto range = partition(arr, l, r, pivot);
        if(idx < range[0]){
            return findTheIndexNum(arr, l, range[0] - 1, idx);
        }else if(idx > range[1]){
            return findTheIndexNum(arr, range[1] + 1, r, idx);
        }
        return arr[l];
    }

public:
    vector<int> smallestK(vector<int>& arr, int k) {
        // 在arr中，找到（如果排序的话）位于k-1位置（下标）的数：
        int kth = findTheIndexNum(arr, 0, arr.size() - 1, k-1); // O(n)
        // 遍历数组，将小于kth的数，全部加入ans
        std::vector<int> ans(k);
        int i = 0;
        for(auto num : arr){
            if(num < kth){
                arr[i++] = num;
            }
        }
        // 如果没选够k个，再遍历一遍数组，将等于kth的数，加入ans，直至选足k个：
        if (k - i > 0) {
            for (int num : arr) {
                if (num == kth) {
                    ans[i++] = num;
                    if (i == k) break;
                }
            }
        }
        return ans;
    }
};

---

我们知道快排每次都会将小于等于基准值的值放到左边，将大于基准值的值放到右边。

因此，我们可以通过判断基准点的下标idx和k的关系来确定过程是否结束。

• idx < k，基准点左侧不足k个，递归处理右边，让基准点下标右移
• idx>k：基准点左侧超过 k 个，递归处理左边，让基准点下标左移；
• idx=k：基准点左侧恰好 k 个，输出基准点左侧元素。