Python：从堆中删除元素

Question

Python 有

heapq

模块，它实现了堆数据结构，并且支持一些基本操作（push、pop）。

如何在 O(log n) 时间内从堆中删除第 i 个元素？是否可以使用

heapq

或者我必须使用另一个模块？

注意，文档底部有一个示例： http://docs.python.org/library/heapq.html 这提出了一种可能的方法 - 这不是我想要的。我希望删除该元素，而不仅仅是将其标记为已删除。

Answer 1

您可以轻松地从堆中删除第 i 个元素：

h[i] = h[-1]
h.pop()
heapq.heapify(h)

只需将要删除的元素替换为最后一个元素，然后删除最后一个元素，然后重新堆化堆即可。这是 O(n)，如果你愿意，你可以在 O(log(n)) 中做同样的事情，但是你需要调用几个内部 heapify 函数，或者更好，正如 larsmans 指出的那样，只需复制将 heapq.py 中的 _siftup/_siftdown 放入您自己的代码中：

h[i] = h[-1]
h.pop()
if i < len(h):
    heapq._siftup(h, i)
    heapq._siftdown(h, 0, i)

请注意，在每种情况下，您不能只执行

h[i] = h.pop()

，因为如果

i

引用最后一个元素，则会失败。如果您在特殊情况下删除最后一个元素，那么您可以结合覆盖和弹出。

请注意，根据堆的典型大小，您可能会发现仅调用

heapify

虽然理论上效率较低，但可能比重新使用

_siftup

/

_siftdown

更快：稍微反省一下就会发现

heapify

可能是用 C 实现的，但内部函数的 C 实现并未公开。如果性能对您很重要，那么请考虑对典型数据进行一些时序测试，看看哪个是最好的。除非你有非常大的堆，否则大 O 可能不是最重要的因素。

编辑：有人试图编辑此答案以删除对

_siftdown

的呼叫，并发表评论：

不需要

_siftdown 。新的 h[i] 保证是旧 h[i] 的子级中最小的，但仍然大于旧 h[i] 的父级 （新 h[i] 的父级）。 _siftdown 将是一个空操作。我必须编辑，因为我 还没有足够的代表来添加评论。

他们在此评论中错过的是

h[-1]

可能根本不是

h[i]

的子项。在

h[i]

插入的新值可能来自堆的完全不同的分支，因此可能需要在任一方向上进行筛选。

还有评论询问为什么不直接使用

sort()

来恢复堆：调用

_siftup

和

_siftdown

都是 O(log n) 操作，调用 heapify 是 O(n) 操作。调用

sort()

是一个 O(n log n) 操作。调用排序很可能足够快，但对于大堆来说，这是不必要的开销。

编辑以避免@Seth Bruder 指出的问题。当

i

引用结束元素时，

_siftup()

调用将会失败，但在这种情况下，从堆末尾弹出元素不会破坏堆不变量。

Answer 2

(a) 考虑一下为什么您不想延迟删除。在很多情况下，这是正确的解决方案。

(b) 堆是一个列表。您可以通过索引删除元素，就像任何其他列表一样，但随后您将需要重新堆化它，因为它将不再满足堆不变量。

Answer 3

实现了一个堆的OOP示例，支持按索引删除。本质上与接受的答案相同，只是 happypath 类示例。

希望这对以后查看这里的人有用。

class RemoveByIndexHeap:
    def __init__(self):
        self.heap = []

    def _swap(self, i: int, j: int):
        self.heap[i], self.heap[j] = self.heap[j], self.heap[i]

    def _sink(self, i: int):
        while i < self.size():
            swap_with = i
            if i * 2 + 1 < self.size() and self.heap[swap_with] > self.heap[i * 2 + 1]:
                swap_with = i * 2 + 1
            if i * 2 + 2 < self.size() and self.heap[swap_with] > self.heap[i * 2 + 2]:
                swap_with = i * 2 + 2
            if swap_with == i:
                break
            else:
                self._swap(i, swap_with)
                i = swap_with

    def _swim(self, i: int):
        while i > 0:
            swap_with = i
            if (i - 1) // 2 >= 0 and self.heap[swap_with] < self.heap[(i - 1) // 2]:
                swap_with = (i - 1) // 2
            if swap_with == i:
                break
            else:
                self._swap(i, swap_with)
                i = swap_with

    def add(self, obj):
        self.heap.append(obj)
        self._swim(self.size() - 1)

    def remove(self, index: int):
        self._swap(index, self.size() - 1)
        self.heap.pop()

        if index != self.size():
            self._sink(index)
            self._swim(index)

    def get_top(self):
        if not self.heap:
            return None

        return self.heap[0]

    def size(self):
        return len(self.heap)