基本思路:寻找最大/最小的k个数,用优先队列(堆)做。因为每次堆超出k个数时需要把堆顶,即最大的那个数pop,所以需要维护大顶堆。

class Solution {
public:
    vector<int> GetLeastNumbers_Solution(vector<int> input, int k) {
        priority_queue<int, vector<int>, less<int>> pq;
        vector<int> res;
        for (auto &i: input) {
            pq.push(i);
            if (pq.size() > k) pq.pop();
        }
        while (!pq.empty()) {
            res.push_back(pq.top());
            pq.pop();
        }
        return res;
    }
};

优化思路:快速排序的思想是选择一个基准元素,每次用O(n)的时间将所有小于基准元素的元素与所有大于基准元素的元素分开, 然后再对两边的元素递归排序。但是我们需要选择的只是k个最小元素,设所有小于基准元素的元素数目为n。

  • 如果n=k则正好满足条件
  • 如果n>k,只需要在这n个元素中递归选择即可
  • 如果n<k,除了这n个元素外还需要再=在大于基准元素的元素中选择k-n个元素,所以只需要在另外一堆元素中递归选择

这样可以做到常数空间复杂度(如果不计算递归消耗的栈),并且由于每个元素至多只会被遍历一次,所以可以做到线性时间复杂度。

class Solution {
public:
    void fastSort(vector<int>& input, int k, int l, int r) {
        if (l >= r) return;
        int base = input[l];
        int sp = l + 1;
        for (int i = l + 1; i <= r; i++) {
            if (input[i] <= base) {
                swap(input[sp++], input[i]);
            }
        }
        swap(input[sp - 1], input[l]);
        if (sp == k) return;
        else if (sp > k) fastSort(input, k, l, sp - 2);
        else fastSort(input, k, sp, r);
    }

    vector<int> GetLeastNumbers_Solution(vector<int> input, int k) {
        fastSort(input, k, 0, input.size() - 1);
        return vector<int>(input.begin(), input.begin() + k);
    }
};