劍指offer：最小的K個數（Python）

題目描述

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

解題思路

不用動腦的解題思路：實現一個排序算法，然後返回前k個數字。這種思路很low，但是考慮到可以幫自己溫習一下快排和歸併排序算法，就先用這兩種排序算法實現了下：

#快速排序，O(nLogn)
def GetLeastNumbers_Solution(self, tinput, k):
    if len(tinput) < k:
        return []
    return self.quick_sort(tinput)[:k]

def quick_sort(self, list):
    if len(list) < 2:
        return list[:]
    left = (self.quick_sort([i for i in list[1:] if i <= list[0]]))
    right = (self.quick_sort([i for i in list[1:] if i > list[0]]))
    return left + [list[0]] + right

#歸併排序，O(nLogn)
def GetLeastNumbers_Solution(self, tinput, k):
    if len(tinput) < k:
        return []
    return self.merge_sort(tinput)[:k]

def merge_sort(self, list):
    if len(list) < 2:
        return list[:]
    left = self.merge_sort(list[:len(list)//2])
    right = self.merge_sort(list[len(list)//2:])
    return self.merge(left,right)

def merge(self,left, right):
    res = []
    while left and right:
        res.append(left.pop(0)) if left[0] < right[0] else res.append(right.pop(0))
    res += left if not right else right
    return res

基於“先整體排序，然後取前k個元素”的思路引申：需要的是前k小的數，沒必要捨近求遠把整個數組都排好序。這樣一來，時間複雜度可以降到O(nLogk)。利用最大堆算法，只不過在這裏稍稍將最大堆算法變爲最小堆算法，代碼如下：

# 最大堆，O(nLogk)
def GetLeastNumbers_Solution(self, tinput, k):
    if len(tinput) < k:
        return []
    res = []
    for i in tinput:
        heapq.heappush(res, -i) if len(res) < k else heapq.heappushpop(res, -i)
    return sorted(list(map(lambda x: -x, res)))

番外：Python中的heapq模塊用來建立“堆”這種數據結構。
heapq.heappush(res, -i) 意爲：向堆res中添加一個元素-i
heapq.heappushpop(res, -i)意爲：將元素-i與堆頂的元素比較。如果該元素值大於堆頂元素，則將該元素與堆頂元素替換。否則不改變堆元素。