可动态改变优先级的优先队列

背景：最近做一个基于图割的图像分割任务，初始将每个像素视为一个独立的region，然后迭代的将最最相似的两个region A 和 B 进行合并，合并后(A+B)到别的邻居的相似度可能会进一步更新。我考虑使用优先队列存储相邻region的相似度，但问题来了，一般的优先队列仅仅支持: 1) 入队：插入元素; 2) 出队：拿到最小元素。但实际情况是，我们需要动态更新region之间的相似度。因此，我们需要修改优先队列，使得其高效的支持修改队列内部元素的操作。

一、实现方式

优先队列使用数组组织数据，实现了O(logN)的入队和出队操作。欲修改某个元素的优先级，那么先得找到之。那么对于一个数组存储的有点队列，查找的时间复杂度是O(N)。 我们对于这样的时间复杂度显然是不能满足的。

我们注意到，优先队列以数组方式存储，内部主要有两个操作：ShiftUp 和 ShiftDown。

ShiftUp：上浮操作，将优先级较低的节点上浮，一般用于队尾插入新的节点后的调整

ShiftDown：下沉操作，将优先级较高的节点下沉，用于出队后队尾换到队头时的调整

其实，无论上浮还是下沉，元素都是在一维数组里互换位置，如下图所示。那么我们可以把我们要存的数据，跟他附加一个key，然后让这个key跟其所在的位置相关联就可以做到。key + data 构成队列的节点，key + position 构成一个map追踪节点的位置变化，实现快速的搜索。

在ShiftUp 和 ShiftDown两个操作中，除了更新数据本身以外，需要实时更新数据节点位置的变化，将其更新到(key, position)MAP中即可。

二、Show me the code

注释写得很详细，就不再解释了。代码参考：https://github.com/kengran/HashHeap/edit/master/main.py，感谢原作者，同时我修正了原始代码的一些错误。

class HashHeap():
    def __init__(self):
        self.A = []  # 真实数据存储在一维数组中，每个节点[key, value], value代表优先级
        self.mapping = {}  # 存储每个节点在一维数组中的位置，字典每项：[key, position]

    def top(self):  # 返回队列头部元素 min
        if len(self.A) == 0: return None
        return self.A[0]

    def pop(self):  # 最小元素出队，同时调整队列
        self.A[0], self.A[-1] = self.A[-1], self.A[0]  # 将队头队尾交换
        key = self.A[0][0]
        del self.mapping[key]  # 字典中删除队头节点对应的项
        self.A.pop()  # 删除队尾
        self.siftDown_(0)  # 将新的队头shift down

    def siftUp_(self, i):  # 上浮操作，将优先级较低的节点上浮，用于队尾插入新的节点 或 手动降低某些节点的优先级后的调整
        while 0 < i < len(self.A) and self.A[(i - 1) // 2][1] > self.A[i][1]:  # 如果当前节点优先级 < 父节点优先级
            self.A[(i - 1) // 2], self.A[i] = self.A[i], self.A[(i - 1) // 2]  # 将当前节点与父节点交换
            self.mapping[self.A[i][0]] = i  # 更新A[i]新的位置
            self.mapping[self.A[(i-1)//2][0]] = (i-1)//2 # 更新A[(i-1)/2]新的位置
            i = (i - 1) // 2  # 继续遍历父节点
        return i

    def siftDown_(self, i):  # 下沉操作，将优先级较高的节点下沉，用于队头出队后队尾换到队头时 或 手动升高某些节点优先级后的调整
        while 2 * i + 1 < len(self.A):
            min_child, l_child, r_child = i, 2 * i + 1, 2 * i + 2
            # 确定左右孩子节点中优先级最小的节点
            if l_child < len(self.A) and self.A[l_child][1] < self.A[min_child][1]:
                min_child = l_child
            if r_child < len(self.A) and self.A[r_child][1] < self.A[min_child][1]:
                min_child = r_child
            # 如果当前节点优先级大于孩子节点，则交换，并更新mapping
            if min_child > i:
                self.A[i], self.A[min_child] = self.A[min_child], self.A[i]
                self.mapping[self.A[i][0]] = i   # 更新A[i]新的位置
                self.mapping[self.A[min_child][0]] = min_child  # 更新A[min_child]新的位置
            else:
                break
            i = min_child

        return i

    def push(self, key, val):
        self.A.append([key, val])  # 追加新的节点
        self.mapping[key] = len(self.A) - 1  # 初始放到队尾的位置
        return self.siftUp_(len(self.A) - 1)  # 最终位置根据调整情况确定

    def getVal(self, key):  # 获取某个节点的优先级
        if key not in self.mapping:
            return -1
        else:
            pos = self.mapping[key]  # 先拿到这个节点的位置
            return self.A[pos][1]  # 再获取这个节点的优先级

    def setVal(self, key, val):  # 修改已有节点
        if key not in self.mapping:
            self.push(key, val)
        else:
            pos = self.mapping[key]
            self.A[pos][1] = val
            pos = self.siftUp_(pos)
            pos = self.siftDown_(pos)

    def remove(self, key):  # 删除已有节点
        if key in self.mapping:
            pos = self.mapping[key]
            self.A[pos], self.A[-1] = self.A[-1], self.A[pos]
            self.mapping[self.A[pos][0]] = pos   # 更新A[pos]新的位置
            self.A.pop()
            del self.mapping[key]
            pos = self.siftUp_(pos)
            pos = self.siftDown_(pos)


if __name__ == "__main__":
    priority_queue = HashHeap()
    priority_queue.setVal('a', 4)
    priority_queue.setVal('b', 3)
    priority_queue.setVal('c', 2)

    print(priority_queue.A)
    print(priority_queue.top())

    priority_queue.pop()

    print(priority_queue.A)
    print(priority_queue.top())

以上代码还有一个C++版本：https://github.com/mgordon9/HashHeap

参考资料：

CSDN: 一个可变优先队列（Mutable Priority Queue）的实现

Stackoverflow: How to update elements within a heap? (priority queue)