參考鏈接
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]新聞熱點
疑難解答