求数组中最小的k个数以及海量数据最大堆、multiset解决方案

【题目】

输入n个整数，找出其中最小的K个数。例如输入4,5,1,6,2,7,3,8这8个数字，则最小的4个数字是1,2,3,4,。

【方案一】

主要有两种方案。第一是利用我们熟知的 partition 算法，它是快速排序的核心，相信每个人都会。它可以用来求取数组的任意第 k 大的数，时间复杂度是O(n)。我们不断对数据 partition，当它返回的 index 为第 k-1 是，那么就说明前 k 个数（包括 index对应的数）就是最小的 k 个数了。因为 index 对应数的左侧都比它小，一共 0~k-2 即 k-1 个，加上它自己，就是 k 个了。

代码：

class Solution {
public:
    vector<int> GetLeastNumbers_Solution(vector<int> input, int k) {
        const int size = input.size();
        vector<int> res;
        if(size == 0 || k <= 0 || k > size)
            return res;
        if(k == size)
            return input;
    
        int start = 0;
        int end = size - 1;
        int index = partition(input, start, end);
        while(index != k - 1){ 
            if(index > k - 1)
                end = index - 1;
            else
                start = index + 1;
            index = partition(input, start, end);
        }   

        for(int i=0; i<k; ++i)
            res.push_back(input[i]);
        return res;
    }   
private:
    int partition(vector<int>& arr, int start, int end){ 
        //int index = ( [start, end] (void)  //我试图利用随机法，但是这不是快排，外部输入不能保证end-start!=0，所以可能发生除零异常
        //              {return random()%(end-start)+start;} )(); 
        //std::swap(arr[start], arr[end]);

        int small = start - 1;
        for(int index=start; index<end; ++index){
            if(arr[index] < arr[end]){
                ++small;
                if(small != index)
                    std::swap(arr[small], arr[index]);
            }
        }
        ++small;
        std::swap(arr[small], arr[end]);
        return small;
    }
};

【方案二】

上面的 partition 算法有两个缺点，其一必须修改原数组元素（除非你拷贝出来，那也太蠢了），其二是不能针对海量数据。所以就有了最大堆的解法。我们用数组前 k 个元素建立 k 个节点的最大堆，后续输入如果小于最大堆的最大值，即头部，那么恭喜它，它入选了。然后把当前最大堆头部换成新元素，重新堆化。继续，循环，直到数据输入完毕。最大堆中剩余的 k 个元素即为所求。

代码：

class Solution {
public:
    vector<int> GetLeastNumbers_Solution(vector<int> input, int k) {
        const int size = input.size();
        vector<int> heap;
        if(size == 0 || k <= 0 || k > size)  //错误返回空
            return heap;
        if(k == size)   //大小相等直接返回
            return input;
        
       	heap.resize(k);  //不能用reserve
        for(int i=0; i<k; ++i)   //将前k个分配给堆
            heap[i] = input[i];
        
        for(int i=(k>>1)-1; i>=0; --i)
            sift_down(heap, i, k);   //建堆
        
        for(int i=k; i<size; ++i){   //遍历第[k+1..n],如果小于最大堆顶，就放入堆顶，然后重新堆化
        	if(input[i] < heap[0]){
                heap[0] = input[i];  //放入堆顶
            	sift_down(heap, 0, k);   //堆化
            }
        }
        
        return heap;
    }
private:
    int left_child(const int i){   //得到左孩子下标
        return (i << 1) + 1;
    }
    void sift_down(vector<int>& heap, int i, const int N){
		int tmp = heap[i];  //保存目标堆化元素
        for(int child = -1; left_child(i)<N; i=child){   //i=child
            child = left_child(i);
            if(child != N-1 && heap[child] < heap[child+1]) //找出左右孩子中较大的一个
                ++child;
            if(heap[child] > tmp)  //如果大于目标，那就让孩子节点覆盖自己
                heap[i] = heap[child];
            else
                break;
        }
        heap[i] = tmp; //这句话不能写在上面的break之上，因为有课能i是叶子节点，left_child(i)<N不会进入循环
    }
};

【方案三】

方案三和方案二理论是一样的，只不过使用了 STL 的 multiset，为什么不用 set 是因为 set 元素不可重复。由于 multiset 内部是红黑树，可以自动排序。我们使用 STL 个 greater<int> 让红黑树由大到小排序，红黑树的 begin() （不是一定头结点）就是最大值了。有了最大值剩下的就和方案二几乎一样了，不再赘述。

代码：

class Solution {
public:
    vector<int> GetLeastNumbers_Solution(vector<int> input, int k) {
        const int size = input.size();
        if(k <= 0 || k > size)
            return std::vector<int>();
        
        std::multiset<int, std::greater<int> > least_nums;
        for(auto v : input){
            if(least_nums.size() < k)
                least_nums.insert(v);
            else{
                auto begin = least_nums.begin();
                if(v < *begin){
                    least_nums.erase(begin);
                    least_nums.insert(v);
                }
            }
        }
        
        return std::vector<int>(least_nums.begin(), least_nums.end());
    }
};