窺探算法之美妙——尋找數組中最小的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)

總結

前面兩種方法在數據量較小的時候如果允許改變數組結構可以使用，但是在大數據場景中，同時不改變數組結構，可以使用第三種方法。