窥探算法之美妙——寻找数组中最小的K个数&python中巧用最大堆

原文发表在我的博客主页，转载请注明出处

前言

不论是小算法或者大系统，堆一直是某种场景下程序员比较亲睐的数据结构，而在python中，由于数据结构的极其灵活性，list，tuple, dict在很多情况下可以模拟其他数据结构，Queue库提供了栈和队列，甚至优先队列（和最小堆类似），heapq提供了最小堆，树，链表的指针在python中可以当作最普通的变量，所以python大法好。。。使用python确实可以把程序员从复杂的数据结构中解放开来，重点关注算法。好了言归正传。

题目

前几天看到了一个很经典的算法题目：输入n个整数，找出其中最小的k个数

解决办法

这道题目本身不是很难，而这篇博客更加侧重的是python中的最大堆的使用以及这道题目的解法汇总。

一. 排序

这个思路应该是最简单的，将整个数组排序，然后取出前k个数据就可以了，这个算法的时间复杂度为nlog(n)，这里展示快速排序。代码如下：

def partition(alist, start, end):
    if end <= start:
        return
    base = alist[start]
    index1, index2 = start, end
    while start < end:
        while start < end and alist[end] >= base:
            end -= 1
        alist[start] = alist[end]
        while start < end and alist[start] <= base:
            start += 1
        alist[end] = alist[start]
    alist[start] = base
    partition(alist, index1, start - 1)
    partition(alist, start + 1, index2)

def find_least_k_nums(alist, k):
    length = len(alist)
    if not alist or k <=0 or k > length:
        return None
    start = 0
    end = length - 1
    partition(alist, start, end)
    return alist[:k]

if __name__ == "__main__":
    l = [1, 9, 2, 4, 7, 6, 3]
    min_k = find_least_k_nums(l, 7)
    print min_k

二. 快速排序的思想

这种解法是在第一种解法上面的一种改进，快速排序的思想大家都已经知道，现在我们只需要最小的k个数，所以如果我们在某次快速排序中，选择的基准树的大小刚好是整个数组的第k小的数据，那么在这次排序完成之后，这个基准数之前的数据就是我们需要的（尽管他们并不是有序的），这个方法同样改变了数组，但是可以将时间复杂度压缩到O(n)，话不多说，直接上代码：

def partition(alist, start, end):
    if end <= start:
        return
    base = alist[start]
    index1, index2 = start, end
    while start < end:
        while start < end and alist[end] >= base:
            end -= 1
        alist[start] = alist[end]
        while start < end and alist[start] <= base:
            start += 1
        alist[end] = alist[start]
    alist[start] = base
    return start

def find_least_k_nums(alist, k):
    length = len(alist)
    #if length == k:
    #    return alist
    if not alist or k <=0 or k > length:
        return
    start = 0
    end = length - 1
    index = partition(alist, start, end)
    while index != k:
        if index > k:
            index = partition(alist, start, index - 1)
        elif index < k:
            index = partition(alist, index + 1, end)
    return alist[:k]

if __name__ == "__main__":
    l = [1, 9, 2, 4, 7, 6, 3]
    min_k = find_least_k_nums(l, 6)
    print min_k

三. 最大堆

上面方法虽然要改变数组的结构，在不要求数字顺序的情况下使用可以获得很好的时间复杂度，但是假如数字非常的多，一次性将其载入内存变得不可能或者内存消耗过大，那上面的方法就不再可行，我们可以创建一个大小为K的数据容器来存储最小的K个数，然后遍历整个数组，将每个数字和容器中的最大数进行比较，如果这个数大于容器中的最大值，则继续遍历，否则用这个数字替换掉容器中的最大值。这个方法的理解也十分简单，至于容器的选择，很多人第一反应便是最大堆，但是python中最大堆如何实现呢？我们可以借助实现了最小堆的heapq库，因为在一个数组中，每个数取反，则最大数变成了最小数，整个数字的顺序发生了变化，所以可以给数组的每个数字取反，然后借助最小堆，最后返回结果的时候再取反就可以了，代码如下：

import heapq
def get_least_numbers_big_data(self, alist, k):
    max_heap = []
    length = len(alist)
    if not alist or k <= 0 or k > length:
        return
    k = k - 1
    for ele in alist:
        ele = -ele
        if len(max_heap) <= k:
            heapq.heappush(max_heap, ele)
        else:
            heapq.heappushpop(max_heap, ele)

    return map(lambda x:-x, max_heap)

if __name__ == "__main__":
    l = [1, 9, 2, 4, 7, 6, 3]
    min_k = get_least_numbers_big_data(l, 3)

总结

前面两种方法在数据量较小的时候如果允许改变数组结构可以使用，但是在大数据场景中，同时不改变数组结构，可以使用第三种方法。