python topN max heap，使用heapq还是自己实现？

Question

python 中有 heapq，用于一般用途。我想记录 topN(0~20) 10e7 条记录。

如果使用 heapq，应使用 '-' 将最大值转换为最小值；并记录底部的最小数量，调用 heapq.heappushpop()

我应该使用 heapq 还是自己实现一个堆(可能有问题或效率较低)？

#update

import heapq
class TopN(object):
    """
    v format: (num, value)

    after looking into http://hg.python.org/cpython/file/2.7/Lib/heapq.py, 
    i find heappushpop already optimize, no need bottom value

    feed() can be optimize further, if needed:
        using func object instead of compare len(self.h) each time
    """
    def __init__(self, N):
        self.N = N
        self.h = []        

    def feed(self, v):  
        if len(self.h) < self.N:
            heapq.heappush(self.h, v)
        else:
            heapq.heappushpop(self.h, v)

    def result(self):
        self.h.sort(reverse=True)
        return self.h

def t_topn():
    topn = TopN(10)
    for i in xrange(5):
        topn.feed((i, str(i)))
    res = topn.result()    
    assert sorted(res, reverse=True) == res 

def t_topn_random():
    import random
    topn = TopN(10)
    for i in xrange(100):
        x = random.randint(0, 1e4)
        topn.feed((x, str(x)))
    res = topn.result()    
    assert sorted(res, reverse=True) == res 

if __name__ == '__main__':
    t_topn()
    t_topn_random()

Answer 1

heapq 的唯一问题是它不像 stdlib 中的其他所有内容那样提供 key 函数。 (如果你很好奇为什么，Raymond Hettinger 在 this email 中解释道。他是对的，heapq 不能提供与其他排序函数相同的接口(interface)——但原因不会影响你的用例，其中 key 只是 lambda x: -x。)

通常的解决方法是 decorate-heap-undecorate。也就是说，将您的值的修改版本放入按 key 排序的堆中。通常，这意味着以下之一:

• 存储key(x)而不是x，然后访问unkey(value)而不是value (假设 key 是可逆的)。
• 存储(key(x), x)而不是x，然后访问value[1]。 (这可能会破坏稳定性，但 heapq 并不保证稳定性。)
• 编写实现自定义 __le__ 方法的包装器类，然后存储 Wrapper(x) 而不是 x 并访问 value .value 而不是 value。

在您的情况下，关键功能是可逆的。因此，只需存储 -x，并访问 -value。这与装饰一样微不足道。

尽管如此，无论它多么简单，您都应该编写一个包装器，否则您会在某些时候搞砸它。例如，您可以编写一个 maxheap，将 minheap 包装在 heapq 中，如下所示:

import heapq
def heapify(x):
    for i in range(len(x)):
        x[i] = -x[i]
    heapq.heapify(x)
def heappush(heap, item):
    heapq.heappush(heap, -item)
def heappop(heap):
    return -heapq.heappop(heap)

... 等等你需要的任何其他功能。这可能有点痛苦，但与从头开始实现整个事情相比，工作量要少得多。

当你这样做的时候，你可能想把堆包装在一个面向对象的 API 中，这样你就可以做 heap.push(x) 而不是 heapq.heappush(heap , x) 等

import heapq
class MaxHeap(object):
    def __init__(self, x):
        self.heap = [-e for e in x]
        heapq.heapify(self.heap)
    def push(self, value):
        heapq.heappush(self.heap, -value)
    def pop(self):
        return -heapq.heappop(self.heap)

…

如果您快速浏览一下 ActiveState 上的食谱或 PyPI 上的模块，您应该会发现其他人已经为您完成了大部分工作。

或者，您可以将 heapq 源代码(它是纯 Python)复制并粘贴为 maxheapq.py 并将 cmp_lt 函数替换为它的对面。 (当然，如果你这样做，修改 cmp_lt 以首先采用 key 参数可能同样容易，而且肯定更清晰，并且修改所有其他函数以传递 key——记住它不再普遍适用，因为它不能像往常一样保证 key只调用一次。)

如果你真的想过危险的生活(你不应该)，你甚至可以用猴子修补它:

import heapq
def cmp_gt(x, y):
    return y < x if hasattr(y, '__lt__') else not (x <= y)
heapq.cmp_lt = cmp_gt

但您不想在实际代码中这样做。