"最小的k个数" 方法汇总

解法一

最容易想到的方法是先对元素进行排序，然后取出前k个数，总时间复杂度O(n*logN)。你一定注意到了，当k=1时，上面的算法对整个数组都进行了排序，而原题目只要求最小的k个数，并不需要前k个数有序，也不需要后n-k个数有序。

如何避免n-k个数的排序呢？我们需要部分排序算法，选择排序和交换排序都是不错的选择。把n个数中的前k个数排序出来，复杂度是O(n*k)。

哪一个更好呢？O(nlogN)还是O(nk)？这取决于k的大小。

解法二

在快速排序中，我们用到了一个辅助函数partition。如果基于数组的第k个数字来调整，使得比k个数字小的所有数字都位于数组的左边，比第k个数字大的所有数字都位于数组的右边。这样调整之后，位于数组左边的k个数字就是最小的k个数字（这k个数字不一定是排序的）。基于这种思路的参考代码如下：

vector kSmallNums(vector &nums, int k) {
    vector kNums;

    int start = 0;
    int end = nums.size() - 1;
    int index = partition(nums, start, end);

    while (index != k - 1) {
        if (index > k - 1) {
            end = index - 1;
        } else {
            start = index + 1;
        }
        index = partition(nums, start, end);
    }
    for (int i = 0; i < k; i++) {
        kNums.push_back(nums[i]);
    }
    return kNums;
}

解法三：

在前面一种算法中，我们运用了二分法的思想。在本算法中，我们同样借助快速排序中的partition函数，但采用的是一种"分而治之"的思想。

假设N个数村粗在S中，我们从数组S中随机找出一个元素X，把数组分为两部分S_a和S_b。S_a中的元素小于等于X，S_b中的元素大于X。这时，有两种可能性：

1. S_a中元素的个数小于K，则S_a中所有的数和S_b中最大的k-|S_a|个元素就是数组S中最大的k个数。

2. S_a中元素的个数大于或等于k，则需要返回S_a中最大的k个元素。

参考代码如下：

int partition(vector &nums, int start, int end) {
    int randomIndex = rand() % nums.size(); //随机选择一个数作为分组标准
    swap(nums[0], nums[randomIndex]);       //避免特殊数据下的算法退化
    int x = nums[0];

    while (start < end) {

        while (start < end && nums[end] > x) {
            end--;
        }
        if (start < end) {
            nums[start] = nums[end];
            start++;
        }

        while (start < end && nums[start] <= x) {
            start++;
        }
        if (start < end) {
            nums[end] = nums[start];
            end--;
        }
    }
    nums[start] = x;
    return start;
}

vector kSmallNums(vector nums, int k) {
    if (k <= 0) {
        return vector();
    }
    if (nums.size() <= k) {
        return nums;
    }
    int index = partition(nums, 0, nums.size() - 1);

    vector left = kSmallNums(vector(nums.begin(), nums.begin() + index), k);
    vector right = kSmallNums(vector(nums.begin() + index, nums.end()), k - left.size());
    left.insert(left.end(), right.begin(), right.end());
    return left;
}

解法四

我们已经得到了三种解法，不过这三个解法有个共同之处，就是需要对数据访问多次，那么就有一个问题，如果N很大呢，100亿？这个时候数据就不能全部装入内存，所以要求尽可能少地遍历所有数据。不妨设N>K，前K个数中最小K个数是一种退化的情况，所有K个数就是最小的K个数。如果考虑第K+1个数X呢？如果X比最小的K个数中的最大的数Y大，那么最小的K个数还是保持不变。如果X比Y小，那么最小的K个数应该去掉Y，而包含X。
进一步，我们用容量为K的最大堆来存储最小的K个数。最大堆的堆顶元素就是最小K个数中最大的一个。每次考虑一个数X，如果X比堆顶元素Y大，则不需要更新堆。否则，用X替换堆顶元素Y。在X替换堆顶元素Y之后，X可能破坏最大堆的结构，需要更新堆来维持堆的性质。更新过程花费的时间复杂度O(logK)。在STL中set和multiset都是基于红黑树来实现，我们可以直接拿来用。
参考代码：

vector kSmallNums(vector nums, int k) {
    multiset> s;
    for (int i = 0; i < nums.size() && i < k; i++) {
        s.insert(nums[i]);
    }
    for (int i = k; i < nums.size(); i++) {
        int num = nums[i];
        if (num < *s.begin()) {
            s.erase(s.begin());
            s.insert(num);
        }
    }
    vector res;
    while (!s.empty()) {
        res.push_back(*s.begin());
        s.erase(s.begin());
    }
    return res;
}

总结

基于函数Partition的第二种解法平均时间复杂度O(n*logK)，第三种解法平均时间复杂度O(n)，他们都需要修改输入的数组。最后一种解法没有修改输入数组，我们每次只是从data中读入数字，所有的写操作都是在堆中进行。二是该算法适合海量数据的输入。由于内存的大小是有限的，有可能不能把海量数据一次性全部读入内存。这个时候，可以从辅助内存中每次读入一个数字。由于每种算法各有优缺点，因此因根据不同的场合采用合适的算法。