算法进阶——最小的K个数

题目

给定一个长度为 n 的可能有重复值的数组，找出其中不去重的最小的 k 个数。例如数组元素是4,5,1,6,2,7,3,8这8个数字，则最小的4个数字是1,2,3,4(任意顺序皆可)。

数据范围：0≤k,n≤10000，数组中每个数的大小0≤val≤1000

要求：空间复杂度 O(n) ，时间复杂度 O(nlogk)

示例1

输入：
[4,5,1,6,2,7,3,8],4 
返回值：
[1,2,3,4]
说明：
返回最小的4个数即可，返回[1,3,2,4]也可以

示例2

输入：
[1],0
返回值：
[]

示例3

输入：
[0,1,2,1,2],3
返回值：
[0,1,1]

思路

利用一个容量为k的multiset来记录临时的最小k个数，循环遍历数组，如果multiset的大小小于k，就直接将元素存入multiset中。存够k个数后，每次将当前元素与multiset的最后一个元素比较，如果当前元素较小，则删除multiset的最后一个元素，将当前元素插入multiset中。最后multiset中就是想要的结果。

解答代码

class Solution {
public:
    /**
     * @param input int整型vector 
     * @param k int整型 
     * @return int整型vector
     */
    vector<int> GetLeastNumbers_Solution(vector<int>& input, int k) {
        // write code here
        vector<int> ret{};
        if (k > input.size() || k == 0) {
            return ret;
        }

        multiset<int> tmp;// 用multiset临时记录最小的k个数
        for (auto num : input) {
            if (tmp.size() < k) {
                tmp.insert(num);
                continue;
            }
            
            auto cur_max_it = --tmp.end();// multiset本身是排序的，取最后一个值就是目前的临时记录中的最大值
            if (num < *cur_max_it) {
                // 找到更小值，则删除最大值的元素，将新的值加到multiset中
                tmp.erase(cur_max_it);
                tmp.insert(num);
            }
        }

        for (auto n : tmp) {
            ret.push_back(n);
        }

        return ret;
    }
};