拦截 heapq

Question

我希望使用Python的heapq模块。不过，我需要跟踪每个值被设置为哪个索引。

于是我写了：

class heap(list):       
    def __init__(self,xs):
        super(heap,self).__init__(xs)
        self._index_table = {x:i for i,x in enumerate(self)}

    def __setitem__(self,i,v):
        print(i,v)
        super(heap,self).__setitem__(i,v)
        self._index_table[v] = i

    def append(self,x):
        super(heap,self).append(x)
        self._index_table[x] = len(self)-1

from heapq import heapify, heappush, heappop, _siftdown, _siftup

h = heap([4,3,2,1])
heapify(h)
heappush(h,12)
print(h)
print(h._index_table)

并且这将打印

[1, 3, 2, 4, 12]
{1: 3, 2: 2, 3: 1, 4: 0}

heapify 和 heappush 修改了列表中的条目，从而规避了我的所有尝试捕获所有赋值操作。

为什么会发生这种情况？有没有办法解决这个问题？是否仍然可以使用 heapq 模块并仍然跟踪每个值对应的索引？

编辑：

看起来我的代码有一行 heapify(h)，这是我在原始帖子中没有的。已经修复了这个问题。

- math4tots

---

我没有具体的建议，但看起来 heapify 和 heappush 调用了修改列表的方法，而你没有重载它们。你尝试过查找 heapq 模块的源代码，看看它调用了哪些列表方法吗？ - Sam Mussmann

@SamMussmann heapq.py 看起来主要使用简单的赋值和追加。但是我刚刚注意到在底部它导入了一个C实现。Python代码本身看起来应该与我的代码兼容... - math4tots

Answer 1

在阅读heapq源代码时，我注意到和@math4tots一样，它正在导入一个C实现。因此，我运行了以下代码来证明它是使用Python源代码（这将调用可以重载的list方法），还是使用使用编译方法的C实现：

>>> class heap(list):        
...     def __init__(self,xs):
...         super(heap,self).__init__(xs)
...         self._index_table = {x:i for i,x in enumerate(self)}
...     
...     def __setitem__(self,i,v):
...         print("SETITEM")
...         print(i,v)
...         super(heap,self).__setitem__(i,v)
...         self._index_table[v] = i
...     
...     def append(self,x):
...         print("APPEND")
...         super(heap,self).append(x)
...         self._index_table[x] = len(self)-1
... 
>>> 
>>> 
>>> h = heap([4,3,2,1])
>>> heapify(h)
>>> h
[1, 3, 2, 4]
>>> h._index_table
{1: 3, 2: 2, 3: 1, 4: 0}
>>> heappush(h,42)
>>> h
[1, 3, 2, 4, 42]
>>> h._index_table
{1: 3, 2: 2, 3: 1, 4: 0}

它没有打印出单个字符串...这意味着它没有使用我们查看的Python源代码，但肯定是编译版本。

因此，您的代码不太可能按原样工作...

阅读heapq模块的C源代码证明了我们的观点：_siftup函数使用PyList_SET_ITEM()从列表中获取值，覆盖任何尝试重载该方法的操作。

尽管如此，希望并未完全破灭。进一步阅读C源代码后，我发现实际上C模块未导出实现heapq算法的_sitf*函数。所以我进行了以下双重检查：

>>> heapq.heapify
<built-in function heapify>
>>> heapq._siftup
<function _siftup at 0x10b36ab00>

这证明了我的正确性！

因此，您可以始终重新实现 heapify() 和 heappush() 函数，这些函数只有几行代码，通过使用 heapq 模块中未被 C 代码遮盖的 _siftup() 和 _siftdown() 函数来实现。

因此，下面是重新实现的方法：

import heapq

class HeapQueue(list):
    def __init__(self,xs):
        super(HeapQueue,self).__init__(xs)
        self._index_table = {x:i for i,x in enumerate(self)}
    def __setitem__(self,i,v):
        super(HeapQueue,self).__setitem__(i,v)
        self._index_table[v] = i
    def append(self,x):
        super(HeapQueue,self).append(x)
        self._index_table[x] = len(self)-1
    def push(self, x):
        self.append(x)
        heapq._siftdown(self, 0, len(self)-1)
    def heapify(self):
        n = len(self)
        for i in reversed(range(n//2)):
            heapq._siftup(self, i)

结果：

>>> h = HeapQueue([4,3,2,1])
>>> print(h._index_table)
{1: 3, 2: 2, 3: 1, 4: 0}
>>> h.heapify()
>>> print(h)
[1, 3, 2, 4]
>>> print(h._index_table)
{1: 0, 2: 2, 3: 1, 4: 3}
>>> h.push(42)
>>> print(h._index_table)
{1: 0, 2: 2, 3: 1, 4: 3, 42: 4}
>>> print(h)
[1, 3, 2, 4, 42]
>>> 

我的猜测是，您不想保留那个heapify()方法，而是将其作为构造函数的一部分，并为自己的堆队列类思考一个更好的接口。我只是为了证明这个想法而已。希望对你有所帮助。

- zmo