面试题40

输入整数数组 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 <= 1000

• 0 <= arr[i] <= 1000

1. heap/priority queue O(nlog(k)) S(k)

2. max heap

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

        for (int i = 0; i < k && i < arr.size(); i++)
            pq.push(arr[i]);

        for (int i = k; i < arr.size(); i++) {
            pq.push(arr[i]);
            pq.pop();
        }

        vector<int> res(pq.size());
        for (int i = res.size() - 1; i >= 0; i--) {
            res[i] = pq.top();
            pq.pop();
        }

        return res;
    }
};

1. quick sort O(n)

2. Time complexity: In average: O(n) + O(n / 2) + O(n / 4) + .... == O(n)

#include <experimental/random>
class Solution {
public:
    vector<int> getLeastNumbers(vector<int>& arr, int k) {
        int n = arr.size(); if (n < k) return arr;
        int lo = 0, hi = n - 1;
        k = k - 1;
        while (lo < hi) {
            swap(arr[lo], arr[std::experimental::randint(lo, hi)]);
            int pivot = arr[lo];
            int i = lo, j = hi;
            while (i < j) {
                while (i < j && arr[j] >= pivot) j--;
                arr[i] = arr[j];
                while (i < j && arr[i] <= pivot) i++;
                arr[j] = arr[i];
            }
            arr[i] = pivot;
            if (i >= k)
                hi = --i;
            if (i <= k)
                lo = ++i;
        }

        return {arr.begin(), arr.begin() + k + 1};
    }
};