数据结构 | 堆 | 问题8

问题描述

设计一个数据结构，支持以下操作：

• push(val)：将元素 val 推入数据结构。
• pop()：删除并返回最大值。
• popAt(idx)：删除并返回在指定下标 idx 的堆中最大值，其中 0 ≤ idx ≤ 总堆数 - 1。

解决方案

此问题可以使用堆和哈希表来解决。使用一个哈希表 hashmap 存储每个元素的下标，使用一个堆 maxheap 存储元素值。

哈希表 hashmap 的键为元素值，值为一个数组，存储该元素所有的下标。堆 maxheap 存储所有元素值，保持最大值在堆顶。

• 对于 push(val) 操作，我们将元素 val 推入 maxheap 中，并将其下标 idx 添加到 hashmap 中对应的值中。
• 对于 pop() 操作，我们从 maxheap 中弹出堆顶，然后删除 hashmap 中对应的元素下标。
• 对于 popAt(idx) 操作，我们获取堆中在指定下标 idx 的元素下标，然后从 hashmap 中将其删除，取出该元素在 maxheap 中对应的元素值，并将其下标从 maxheap 中删除。

代码如下：

import heapq
from collections import defaultdict

class MaxHeap:
    def __init__(self):
        self.heap = []
        self.count = 0

    def push(self, val):
        self.count += 1
        heapq.heappush(self.heap, -val)

    def pop(self):
        self.count -= 1
        return -heapq.heappop(self.heap)

    def popAt(self, idx):
        if idx >= self.count:
            return None
        heap = self.heap
        val = -heap[idx]
        heap[idx] = -float('inf')
        heapq.heapify(heap)
        self.count -= 1
        return val

class MaxHeapHashMap:
    def __init__(self):
        self.maxheap = MaxHeap()
        self.hashmap = defaultdict(list)

    def push(self, val):
        self.maxheap.push(val)
        self.hashmap[val].append(self.maxheap.count-1)

    def pop(self):
        val = self.maxheap.pop()
        idx = self.hashmap[val][-1]
        self.hashmap[val].pop()
        if not self.hashmap[val]:
            del self.hashmap[val]
        return val

    def popAt(self, idx):
        if idx >= self.maxheap.count:
            return None
        val = self.maxheap.popAt(idx)
        if val is None:
            return None
        self.hashmap[val].remove(idx)
        if not self.hashmap[val]:
            del self.hashmap[val]
        return val

性能分析

时间复杂度

• push(val)：时间复杂度为 $O(logn)$，其中 $n$ 是堆的大小。
• pop()：时间复杂度为 $O(logn)$，其中 $n$ 是堆的大小。
• popAt(idx)：时间复杂度为 $O(logn)$，其中 $n$ 是堆的大小。

空间复杂度

空间复杂度为 $O(n)$，其中 $n$ 是堆的大小。哈希表和堆都需要存储 $n$ 个元素。