Python中如何实现O(log n)时间删除二叉堆中非根元素？

首先直接给结论：是的，确实存在可以在O(log n)时间内删除任意非根元素的二叉堆实现，但Python标准库的heapq模块本身并没有提供这个功能——这也是你遇到heap.index(wKeys)速度极慢的核心原因，因为这个方法是线性遍历整个堆（O(n)时间）来定位元素，完全拖慢了你的操作效率。

接下来给你几个实用的解决方案，按实现成本从低到高排序：

1. 用现成的平衡二叉搜索树库（最省心） #

如果你不想自己折腾堆的实现，强烈推荐使用SortedList（来自sortedcontainers第三方库）。它本质是一个平衡二叉搜索树，既能像堆一样快速获取最小/最大元素（O(1)时间），又支持O(log n)时间的任意元素删除，完美匹配你的业务场景。

举个简单的使用示例：

from sortedcontainers import SortedList

# 初始化一个"最小堆"
priority_heap = SortedList()

# 添加元素
priority_heap.add(3)
priority_heap.add(1)
priority_heap.add(4)

# 获取堆顶（最小元素）
print(priority_heap[0])  # 输出: 1

# 删除任意非根元素，比如3
priority_heap.discard(3)  # O(log n)时间

# 再次获取堆顶
print(priority_heap[0])  # 输出: 4

这个方案不需要你关心底层数据结构的维护，直接调用API就能完成需求，是最高效的选择。

2. 自定义带索引映射的二叉堆（手动实现） #

如果你坚持要用二叉堆，可以自己实现一个带位置映射的版本。核心思路是维护一个字典，记录每个元素在堆数组中的索引位置，然后通过decrease_key/increase_key操作把要删除的元素调整到堆顶，再用heappop弹出（这两步都是O(log n)时间）。

注意：这个实现需要在堆的siftup/siftdown操作中同步更新位置映射，否则会出现索引失效的问题。这里给你一个简化版的示例：

import heapq

class IndexedMinHeap:
    def __init__(self):
        self.heap = []
        # 映射：元素 -> 在堆中的索引
        self.pos_map = {}
        # 计数器处理重复元素，避免优先级相同的元素比较冲突
        self.counter = 0

    def push(self, priority, item):
        entry = (priority, self.counter, item)
        heapq.heappush(self.heap, entry)
        self.pos_map[item] = len(self.heap) - 1
        self.counter += 1

    def _sift_up(self, idx):
        # 自定义sift_up，同步更新pos_map
        while idx > 0:
            parent_idx = (idx - 1) // 2
            if self.heap[idx] < self.heap[parent_idx]:
                # 交换元素
                self.heap[idx], self.heap[parent_idx] = self.heap[parent_idx], self.heap[idx]
                # 更新映射
                self.pos_map[self.heap[idx][2]] = idx
                self.pos_map[self.heap[parent_idx][2]] = parent_idx
                idx = parent_idx
            else:
                break

    def _sift_down(self, idx):
        # 自定义sift_down，同步更新pos_map
        size = len(self.heap)
        while True:
            left_child = 2 * idx + 1
            right_child = 2 * idx + 2
            smallest = idx

            if left_child < size and self.heap[left_child] < self.heap[smallest]:
                smallest = left_child
            if right_child < size and self.heap[right_child] < self.heap[smallest]:
                smallest = right_child

            if smallest != idx:
                self.heap[idx], self.heap[smallest] = self.heap[smallest], self.heap[idx]
                self.pos_map[self.heap[idx][2]] = idx
                self.pos_map[self.heap[smallest][2]] = smallest
                idx = smallest
            else:
                break

    def delete(self, item):
        if item not in self.pos_map:
            raise ValueError("元素不在堆中")
        
        idx = self.pos_map[item]
        # 将目标元素的优先级设为极小值，调整到堆顶
        self.heap[idx] = (float('-inf'), self.counter, item)
        self._sift_up(idx)
        # 弹出堆顶元素
        self.pop_top()

    def pop_top(self):
        if not self.heap:
            raise IndexError("堆为空")
        priority, _, item = heapq.heappop(self.heap)
        del self.pos_map[item]
        # 如果堆还有元素，更新最后一个元素的映射
        if self.heap:
            self.pos_map[self.heap[0][2]] = 0
        return priority, item

这个实现能做到O(log n)的任意元素删除，但需要你自己维护堆的结构和映射关系，适合对底层实现有要求的场景。

3. 关于Treap的补充 #

你提到的Treap本身是支持O(log n)删除任意元素的，因为它是平衡二叉搜索树的一种。只是你看到的那个实现没有暴露这个接口而已。不过自己实现Treap的成本远高于前面两种方案，所以除非有特殊需求，否则不推荐。

内容的提问来源于stack exchange，提问作者Endre Moen