Python中的最小堆

Question

我想把一组对象存储在一个最小堆里，并且想定义一个自定义的比较函数。我注意到Python里有一个叫heapq的模块。请问这个模块可以使用自定义的比较器吗？如果不行，有没有人自己做过一个自定义的最小堆？

Answer 1

有两个选择（除了Devin Jeanpierre的建议）：

1. 在使用堆之前对你的数据进行装饰。这就像给排序加上一个key=选项。例如，如果你（出于某种原因）想根据数字的正弦值来堆化一个数字列表：

   data = [ # list of numbers ]
   heap = [(math.sin(x), x) for x in data]
   heapq.heapify(heap)
   # get the min element
   item = heappop(heap)[1]
   

2. heapq模块是用纯Python写的。你可以直接把它复制到你的工作目录，然后修改相关的部分。简单看一下，你需要修改siftdown()和siftup()，如果你需要的话，可能还要改一下nlargest和nsmallest。

Answer 2

是的，有办法。你可以定义一个包装类，这个类里面实现你自己的比较器，然后用这个包装类的列表来代替你实际对象的列表。这样做是使用heapq模块时最好的选择，因为heapq模块不像排序函数那样提供key=或cmp=这些参数。

def gen_wrapper(cmp):
    class Wrapper(object):
        def __init__(self, value): self.value = value
        def __cmp__(self, obj): return cmp(self.value, obj.value)
    return Wrapper