如何实现支持O(log n)优先级更新的优先队列？适配Prim算法

我来帮你解决这个问题——要实现O(log n)的优先级更新，核心瓶颈就是快速定位元素在堆中的索引，原代码里的线性查找（O(n)）直接拖慢了整体复杂度。我们可以通过引入一个辅助哈希表，维护元素值到堆索引的映射，轻松实现O(1)的元素位置访问，进而让整个updatePriority操作达到O(log n)的时间复杂度。

具体修改步骤 #

1. 在PriorityQueue中添加哈希表映射 #

首先，我们需要在PriorityQueue类里新增一个HashMap<V, Integer>，用来记录每个元素值对应的堆索引位置：

public class PriorityQueue<K, V> {
    private Heap<K, V> heap;
    private Comparator<? super K> comparator;
    // 新增：用哈希表映射元素值到堆中的索引
    private HashMap<V, Integer> valueToIndexMap;

    public PriorityQueue(Comparator<? super K> comparator) {
        this.comparator = comparator;
        valueToIndexMap = new HashMap<>();
        // 把哈希表传给Heap，让堆操作时同步维护映射
        heap = new Heap<>(comparator, valueToIndexMap);
    }
}

2. 修改Heap类，同步维护哈希表 #

堆在执行swap、add、updatePriority这些操作时，元素的索引会动态变化，所以我们需要让Heap在操作时同步更新哈希表的映射。这里直接把哈希表的引用传给Heap，操作时实时修改：

import java.util.Arrays;
import java.util.NoSuchElementException;
import java.util.HashMap;
import java.util.Comparator;

public class Heap<K, V> {
    private int size;
    private Element<K, V>[] heap;
    private Comparator<? super K> comparator;
    // 持有PriorityQueue的哈希表引用，用于同步索引
    private HashMap<V, Integer> valueToIndexMap;

    @SuppressWarnings("unchecked")
    public Heap(Comparator<? super K> comparator, HashMap<V, Integer> valueToIndexMap) {
        this.comparator = comparator;
        this.valueToIndexMap = valueToIndexMap;
        // 初始化堆数组，默认容量16
        heap = (Element<K, V>[]) new Element[16];
        size = 0;
    }

    // 辅助方法：获取父节点索引（0-based堆）
    private int parent(int i) {
        return (i - 1) / 2;
    }

    // 交换元素时同步更新哈希表
    private void swap(int i, int j) {
        Element<K, V> temp = heap[i];
        heap[i] = heap[j];
        heap[j] = temp;
        // 更新两个元素的索引映射
        valueToIndexMap.put(heap[i].value, i);
        valueToIndexMap.put(heap[j].value, j);
    }

    // 修改add方法，添加元素时同步记录索引
    public void add(K priority, V value) {
        size++;
        // 扩容逻辑：当堆满时，容量翻倍
        if (size > heap.length) {
            heap = Arrays.copyOf(heap, heap.length * 2);
        }
        int i = size - 1; // 0-based索引，原代码的i=size是1-based，这里修正
        heap[i] = new Element<>(priority, value);
        // 记录元素的初始索引
        valueToIndexMap.put(value, i);

        // 上浮操作：如果当前节点优先级高于父节点，交换位置
        while (i > 0 && comparator.compare(heap[parent(i)].priority, heap[i].priority) < 0) {
            swap(i, parent(i));
            i = parent(i);
        }
    }

    // 修改updatePriority方法，直接修改优先级而非新建Element
    public void updatePriority(int i, K newPriority) {
        K oldPriority = heap[i].priority;
        heap[i].priority = newPriority; // 直接修改更高效，无需新建对象

        if (comparator.compare(oldPriority, newPriority) > 0) {
            // 优先级降低，执行下沉操作
            heapify(i);
        } else {
            // 优先级升高，执行上浮操作
            while (i > 0 && comparator.compare(heap[parent(i)].priority, heap[i].priority) < 0) {
                swap(i, parent(i));
                i = parent(i);
            }
        }
    }

    // 补充下沉操作（原代码缺失的heapify）
    private void heapify(int i) {
        int leftChild = 2 * i + 1;
        int rightChild = 2 * i + 2;
        int largest = i;

        // 找到当前节点、左孩子、右孩子中优先级最高的
        if (leftChild < size && comparator.compare(heap[leftChild].priority, heap[largest].priority) > 0) {
            largest = leftChild;
        }
        if (rightChild < size && comparator.compare(heap[rightChild].priority, heap[largest].priority) > 0) {
            largest = rightChild;
        }

        // 如果最高优先级不是当前节点，交换并递归下沉
        if (largest != i) {
            swap(i, largest);
            heapify(largest);
        }
    }

    // 可选：Prim算法需要提取堆顶元素，这里补充extractTop方法
    public Element<K, V> extractTop() {
        if (size == 0) {
            throw new NoSuchElementException("Heap is empty");
        }
        Element<K, V> top = heap[0];
        // 从哈希表中移除该元素的映射
        valueToIndexMap.remove(top.value);

        // 将最后一个元素移到堆顶
        heap[0] = heap[size - 1];
        valueToIndexMap.put(heap[0].value, 0);
        size--;

        // 下沉堆顶元素，维护堆结构
        heapify(0);
        return top;
    }
}

3. 实现PriorityQueue的updatePriority方法 #

现在我们可以通过哈希表O(1)获取元素的索引，直接调用Heap的updatePriority即可：

public void updatePriority(K newPriority, V elem) {
    Integer index = valueToIndexMap.get(elem);
    if (index == null) {
        // 元素不在队列中，可根据需求抛出异常或处理
        throw new IllegalArgumentException("Element not present in the priority queue");
    }
    heap.updatePriority(index, newPriority);
}

为什么这个方案可行？ #

• 哈希表的get和put操作平均时间复杂度是O(1)，保证了元素索引的快速获取
• 堆的updatePriority操作依赖上浮/下沉，最多遍历堆的高度（O(log n)）
• 所有堆操作（add、swap、extractTop）都会实时同步哈希表的映射，确保索引的准确性

针对Prim算法的适配说明 #

Prim算法中，每个图节点只会被加入优先队列一次，后续仅更新其对应的权重（比如找到更短的边），不会出现重复元素的情况，所以哈希表的键（元素值V）是唯一的，完全适配这个场景。

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