Java中如何基于集合与自定义比较器O(N)构造PriorityQueue

方案1：反射实现（JDK8环境下可达到O(n)时间复杂度） #

JDK8的PriorityQueue内部已经包含了heapify()的O(n)建堆逻辑，只是没有对外暴露同时传入集合和比较器的构造方法，我们可以通过反射直接操作内部成员完成构造：

import java.lang.reflect.Field;
import java.lang.reflect.Method;
import java.util.Collection;
import java.util.Comparator;
import java.util.PriorityQueue;

public class PriorityQueueUtils {
    @SuppressWarnings("unchecked")
    public static <E> PriorityQueue<E> create(Collection<? extends E> c, Comparator<? super E> comparator) throws Exception {
        // 初始化指定容量、绑定自定义比较器的空队列
        PriorityQueue<E> queue = new PriorityQueue<>(c.size(), comparator);
        // 反射获取内部存储数组
        Field queueField = PriorityQueue.class.getDeclaredField("queue");
        queueField.setAccessible(true);
        Object[] innerQueue = (Object[]) queueField.get(queue);
        // 拷贝集合元素到内部数组
        int size = 0;
        for (E element : c) {
            if (element == null) {
                throw new NullPointerException("优先级队列不支持null元素");
            }
            innerQueue[size++] = element;
        }
        // 反射更新队列size
        Field sizeField = PriorityQueue.class.getDeclaredField("size");
        sizeField.setAccessible(true);
        sizeField.setInt(queue, size);
        // 反射调用heapify方法完成O(n)建堆
        Method heapifyMethod = PriorityQueue.class.getDeclaredMethod("heapify");
        heapifyMethod.setAccessible(true);
        heapifyMethod.invoke(queue);
        return queue;
    }
}

注意：该方案依赖JDK内部实现，不同厂商、不同版本的JDK可能存在字段/方法名变更的问题，对跨版本兼容性要求高的场景建议使用方案2。

方案2：自定义优先级堆实现（无依赖、兼容性强） #

如果不想引入反射，可以自行实现支持自定义比较器、支持O(n)建堆的堆结构，核心逻辑如下：

import java.util.Collection;
import java.util.Comparator;

public class CustomPriorityQueue<E> {
    private final Object[] queue;
    private int size;
    private final Comparator<? super E> comparator;

    public CustomPriorityQueue(Collection<? extends E> c, Comparator<? super E> comparator) {
        this.comparator = comparator;
        this.queue = c.toArray();
        this.size = queue.length;
        // 非空集合执行O(n)堆化
        if (size > 1) {
            heapify();
        }
    }

    @SuppressWarnings("unchecked")
    private void heapify() {
        // 从最后一个非叶子节点开始自顶向下调整
        for (int i = (size >>> 1) - 1; i >= 0; i--) {
            siftDown(i, (E) queue[i]);
        }
    }

    @SuppressWarnings("unchecked")
    private void siftDown(int index, E current) {
        int half = size >>> 1;
        while (index < half) {
            int leftChild = (index << 1) + 1;
            E minChild = (E) queue[leftChild];
            int rightChild = leftChild + 1;
            // 取左右子节点中更小的那个
            if (rightChild < size && comparator.compare(minChild, (E) queue[rightChild]) > 0) {
                minChild = (E) queue[leftChild = rightChild];
            }
            // 当前节点比最小子节点小，调整结束
            if (comparator.compare(current, minChild) <= 0) {
                break;
            }
            // 否则下沉当前节点
            queue[index] = minChild;
            index = leftChild;
        }
        queue[index] = current;
    }

    // 按需补充poll、peek、add、remove等优先级队列常规方法即可
}

补充说明 #

如果可以升级JDK版本，JDK9及以上版本的PriorityQueue已经官方提供了你需要的构造方法，直接调用即可：

// JDK9+ 原生支持同时传入集合和自定义比较器
PriorityQueue<Foo> queue = new PriorityQueue<>(inputList, customComparator);

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