【劍指Offer系列40】最小的k個數

題目

輸入整數數組 arr ，找出其中最小的 k 個數。例如，輸入4、5、1、6、2、7、3、8這8個數字，則最小的4個數字是1、2、3、4。

代碼

Python

class Solution:
    def getLeastNumbers(self, arr: List[int], k: int) -> List[int]:
        if k==0: return list()
        
        hp = [-x for x in arr[:k]] # 準備建立最大堆，負號因爲後面調用堆的方法彈出最小值
        heapq.heapify(hp) # 將列表轉換爲堆
        # 遍歷列表剩餘元素，最後堆中留下的就是那k個最小值
        for i in range(k,len(arr)):
            if -hp[0] > arr[i]: # 堆頂元素大於待插入元素
                heapq.heappop(hp) # 彈出堆內最小值
                heapq.heappush(hp,-arr[i]) # 壓入待插入元素的相反數
        res = [-x for x in hp]
        return res

C++

class Solution {
public:
    vector<int> getLeastNumbers(vector<int>& arr, int k) {
        vector<int> vec(k,0); // 創建一個有k個0的數組
        int n=arr.size();
        if (k==0) return vec;
        
        priority_queue<int> hp; // 堆就是一種優先隊列
        for (int i=0; i<k; ++i) hp.push(arr[i]); // 先建立k個數的堆,採用++i比i++更快
        for (int j=k; j<n; ++j) {
            if (hp.top() > arr[j]) {
                hp.pop();
                hp.push(arr[j]);
            }
        }
        for (int m=0; m<k; ++m) {
            vec[m]=hp.top();
            hp.pop();
        }
        return vec;
    }
};