剑指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与堆顶的元素比较。如果该元素值大于堆顶元素，则将该元素与堆顶元素替换。否则不改变堆元素。