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

總結

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