剑指 Offer 40 最小的k个数

    
        标签：
        
            剑指 Offer
        
    发布于：2022-02-26 14:06:55
    编辑于：2022-02-26 17:17:57
    浏览量：767

概述

https://leetcode-cn.com/problems/zui-xiao-de-kge-shu-lcof/

类似题目，求第 k 大的数： https://leetcode.com/problems/kth-largest-element-in-an-array/

labuladong 也有讲：https://labuladong.gitee.io/algo/4/31/128/

排序法

升序排序后返回前 k 个即可。

class Solution {
public:
    vector<int> getLeastNumbers(vector<int>& arr, int k) {
        sort(arr.begin(), arr.end());
        vector<int> ans;
        k --;
        while (k >= 0) {
            ans.push_back(arr[k--]);
        }
        return ans;
    }
};

O(nlogn) 的时间复杂度。快排需要 O(logn) 的栈空间，因此空间复杂度 O(logn)。

借鉴快排的思想

https://leetcode-cn.com/problems/zui-xiao-de-kge-shu-lcof/solution/zui-xiao-de-kge-shu-by-leetcode-solution/

时间复杂度：期望为 O(n)O(n) ，由于证明过程很繁琐，所以不在这里展开讲。具体证明可以参考《算法导论》第 9 章第 2 小节。最坏情况下的时间复杂度为 O(n^2)。情况最差时，每次的划分点都是最大值或最小值，一共需要划分 n - 1n−1 次，而一次划分需要线性的时间复杂度，所以最坏情况下时间复杂度为 O(n^2)

空间复杂度：期望为 O(logn)，递归调用的期望深度为 O(logn)，每层需要的空间为 O(1)，只有常数个变量。最坏情况下的空间复杂度为 O(n)。最坏情况下需要划分 n 次，即 randomized_selected 函数递归调用最深 n - 1 层，而每层由于需要 O(1) 的空间，所以一共需要 O(n) 的空间复杂度。

class Solution {
public:
    vector<int> getLeastNumbers(vector<int>& arr, int k) {
        helper(arr, k, 0, arr.size() - 1);
        vector<int> ans;
        for (int i = 0; i < k; i ++) ans.push_back(arr[i]);
        return ans;
    }

    void helper(vector<int>& arr, int k, int l, int r) {
        if (l >= r) return;
        int p = partition(arr, l, r);
        if (p == k) {
            return;
        } else if (p > k) {
            helper(arr, k, l, p - 1);
        } else {
            helper(arr, k, p + 1, r);
        }
    }

    int partition(vector<int>& arr, int l, int r) {
        int p = arr[r];
        int i = l - 1;
        for (int j = l; j <= r - 1; j ++) {
            if (arr[j] <= p) {
                i ++;
                swap(arr[i], arr[j]);
            }
        }
        swap(arr[i+1], arr[r]);
        return i + 1;
    }
};

最小堆

最小堆，压入所有元素后弹出 k 个即可。

class Solution {
public:
    vector<int> getLeastNumbers(vector<int>& arr, int k) {
        priority_queue<int, vector<int>, greater<int>> minHeap;
        for (auto n : arr) minHeap.push(n);
        vector<int> ans;
        for (int i = 0; i < k; i ++) {
            ans.push_back(minHeap.top());
            minHeap.pop();
        }
        return ans;
    }
};

时间复杂度：O(nlogn)，因为我们把 n 个元素都插入了，每次插入 O(logn)

空间复杂度：O(n)

最大堆

实际上用最大堆也是可以的哦，而且更高效。

我们这里用了 top()，需要检查 heap 是否为空，其仅在 k 为 0 时可以为空，所有需要对 k == 0 的情况进行特殊处理。

class Solution {
public:
    vector<int> getLeastNumbers(vector<int>& arr, int k) {
        vector<int> ans;
        if (k == 0) return ans;
        priority_queue<int, vector<int>> maxHeap;
        for (int i = 0; i < k; i ++) {
            maxHeap.push(arr[i]);
        }
        for (int i = k; i < arr.size(); i ++) {
            if (arr[i] < maxHeap.top()) {
                maxHeap.pop();
                maxHeap.push(arr[i]);
            }
        }
        while (!maxHeap.empty()) {
            ans.push_back(maxHeap.top());
            maxHeap.pop();
        }
        return ans;
    }
};

时间复杂度：O(nlogk)，最坏情况下我们插入 n 个元素，每次插入 O(logk)

空间复杂度：堆中最多有 k 个元素，所以是 O(logk)

Links: sword-offer-40