python使用heapq實現小頂堆（TopK大）/大頂堆（BtmK小）

參考鏈接

https://www.coder4.com/archives/3844

求一個數列前K大數的問題經常會遇到，在程序中一般用小頂堆可以解決，下面的代碼是使用python的heapq實現的小頂堆示例代碼：

    # !/usr/bin/env python
    # -*- coding:gbk -*-

    import sys
    import heapq

    class TopKHeap(object):
        def __init__(self, k):
            self.k = k
            self.data = []

        def push(self, elem):
            if len(self.data) < self.k:
                heapq.heappush(self.data, elem)
            else:
                topk_small = self.data[0]
                if elem > topk_small:
                    heapq.heapreplace(self.data, elem)
        def topk(self):
            return [x for x in reversed([heapq.heappop(self.data) for x in xrange(len(self.data))])]


    def main():
        list_num = [1, 2, 3, 4, 5, 6, 7, 8, 9]
        th = TopKHeap(5)

        for i in list_num:
            th.push(i)

        print th.topk()

    if __name__ == "__main__":
        main()

python的heapq在實現的時候，沒有像STL或者Java可以傳入比較函數，具體的原因可以參考參考文檔給出的鏈接。

因此有些人想出了比較trick的思路。一句話概括如下：

push(e)改爲push(-e)，pop(e)爲-pop(e)，也就是說存入和取出的數都是相反數，其他邏輯和TopK相同。（點贊）

實現用戶自定義的比較函數，允許elem是一個tuple，按照tuple的第一個元素進行比較，所以可以把tuple的第一個元素作爲我們的比較的key。

英文原話：

The heapq documentation suggests that heap elements could be tuples in which the first element is the priority and defines the sort order.

    import heapq
    class MyHeap(object):
       def __init__(self, initial=None, key=lambda x:x):
           self.key = key
           if initial:
               self._data = [(key(item), item) for item in initial]
               heapq.heapify(self._data)
           else:
               self._data = []

       def push(self, item):
           heapq.heappush(self._data, (self.key(item), item))

       def pop(self):
           return heapq.heappop(self._data)[1]