【題目】

輸入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());
    }
};