PriorityQueue 简介

手写一个 PriorityQueue

二叉堆简介

基于 JDK 动态数组实现二叉堆

源码实现

/**
 * 元素是可比较的小顶堆
 */
public class MinHeap<E> {

    private List<E> list;
    private Comparator<E> comparator;
    // ...
}

/**
 * 创建一个空的最小堆
 */
public MinHeap() {
    list = new ArrayList<>();
}

/**
 * 创建一个空的最小堆，并指定比较器来决定元素之间的顺序
 */
public MinHeap(Comparator<E> comparator) {
    list = new ArrayList<>();
    this.comparator = comparator;
}

/**
 * 参数为数组，将给定的数组转换为一个最小堆
 */
public MinHeap(E[] arr) {
    list = new ArrayList<>(Arrays.asList(arr));
    heapify();
}

/**
 * 创建指定容量的最小堆
 */
public MinHeap(int capacity) {
    list = new ArrayList<>(capacity);
}


/**
 * 创建指定容量的最小堆并指定比较器
 */
public MinHeap(int capacity, Comparator<E> comparator) {
    list = new ArrayList<>(capacity);
    this.comparator = comparator;
}

/**
 * 返回堆中元素个数
 */
public int size() {
    return list.size();
}

/**
 * 判断堆元素是否为空
 */
public boolean isEmpty() {
    return list.isEmpty();
}


/**
 * 获取当前节点的父节点索引
 */
private int parentIndex(int childIndex) {
    if (childIndex == 0) {
        throw new IllegalArgumentException();
    }

    return (childIndex - 1) / 2;
}

/**
 * 返回当前节点的左子节点
 */
private int leftIndex(int index) {
    return 2 * index + 1;
}

/**
 * 返回当前节点的右子节点
 */
private int rightIndex(int index) {
    return 2 * index + 2;
}

/**
 * 将元素存到小顶堆中
 */
public void add(E e) {
    list.add(e);
    siftUp(list.size() - 1);
}

/**
 * 为了保证新增节点后，数组仍符合小顶堆特性，这里需要siftUp保持一下平衡
 */
private void siftUp(int index) {
    if (comparator != null) {
        //循环自新增节点开始自底向上比较子节点和父节点大小，如果子节点大于父节点则交换两者的值
        while (index > 0 && comparator.compare(list.get(parentIndex(index)), list.get(index)) > 0) {
            E tmpElement = list.get(index);
            list.set(index, list.get(parentIndex(index)));
            list.set(parentIndex(index), tmpElement);
            index = parentIndex(index);
        }
    } else {
        //循环自新增节点开始自底向上比较子节点和父节点大小，如果子节点大于父节点则交换两者的值
        while (index > 0 && ((Comparable) list.get(parentIndex(index))).compareTo(list.get(index)) > 0) {
            E tmpElement = list.get(index);
            list.set(index, list.get(parentIndex(index)));
            list.set(parentIndex(index), tmpElement);
            index = parentIndex(index);
        }
    }
}

/**
 * 获取小顶堆堆顶的元素
 */
public E poll() {
    if (list == null || list.isEmpty()) {
        return null;
    }
    E ret = list.get(0);
    list.set(0, list.get(list.size() - 1));
    list.remove(list.size() - 1);
    siftDown(0);
    return ret;
}

/**
 * 将数组转为堆
 */
private void heapify() {
    //找到最后一个非叶子节点,往前遍历
    for (int i = parentIndex(list.size() - 1); i >= 0; i--) {
        //对每个非叶子节点执行siftDown,使其范围内保持小顶堆特性
        siftDown(i);
    }
}    

private void siftDown(int index) {

    if (comparator != null) {
        //如果左节点小于数组大小才进入循环，避免数组越界
        while (leftIndex(index) < list.size()) {
            //获取左索引
            int cmpIdx = leftIndex(index);

            //获取左或者右子节点中值更小的索引
            if (rightIndex(index) < list.size() &&
                    comparator.compare(list.get(leftIndex(index)), list.get(rightIndex(index))) > 0) {
                cmpIdx = rightIndex(index);
            }

            //如果父节点比子节点小则停止比较，结束循环
            if (comparator.compare(list.get(index), list.get(cmpIdx)) <= 0) {
                break;
            }

            //反之交换位置，将index更新为交换后的索引index，进入下一个循环
            E tmpElement = list.get(cmpIdx);
            list.set(cmpIdx, list.get(index));
            list.set(index, tmpElement);


            index = cmpIdx;

        }
    } else {
        //如果左节点小于数组大小才进入循环，避免数组越界
        while (leftIndex(index) < list.size()) {
            //获取左索引
            int cmpIdx = leftIndex(index);

            //获取左或者右子节点中值更小的索引
            if (rightIndex(index) < list.size() &&
                    ((Comparable) list.get(leftIndex(index))).compareTo(list.get(rightIndex(index))) > 0) {
                cmpIdx = rightIndex(index);
            }

            //如果父节点比子节点小则停止比较，结束循环
            if (((Comparable) list.get(index)).compareTo(list.get(cmpIdx)) <= 0) {
                break;
            }

            //反之交换位置，将index更新为交换后的索引index，进入下一个循环
            E tmpElement = list.get(cmpIdx);
            list.set(cmpIdx, list.get(index));
            list.set(index, tmpElement);
            index = cmpIdx;

        }
    }


}

int parentIndex=(childIndex - 1) / 2

/**
 * 将数组转为堆
 */
private void heapify() {
    //找到最后一个非叶子节点,往前遍历
    for (int i = parentIndex(list.size() - 1); i >= 0; i--) {
        //对每个非叶子节点执行siftDown,使其范围内保持小顶堆特性
        siftDown(i);
    }
}

public E peek() {
    if (list == null || list.isEmpty()) {
        return null;
    }
    return list.get(0);
}

测试

int n = 1000_0000;

MinHeap<Integer> heap = new MinHeap<>(n, Comparator.comparingInt(a -> a));

//往堆中随机存放1000w个元素
for (int i = 0; i < n; i++) {
    heap.add(ThreadLocalRandom.current().nextInt(0, n));
}

int[] arr = new int[n];

//进行出队操作，并将元素存到数组中
for (int i = 0; i < n; i++) {
    arr[i] = heap.poll();
}

//循环遍历，如果前一个元素比后一个元素大，则说明我们的小顶堆有问题
for (int i = 1; i < n; i++) {
    if (arr[i - 1] > arr[i]) {
        throw new RuntimeException("err heap");
    }
}

小结

基于二叉堆实现一个 PriorityQueue

public interface Queue<E> {

    /**
     * 获取当前队列大小
     */
    int size();

    /**
     * 判断当前队列是否为空
     */
    boolean isEmpty();

    /**
     * 添加元素到优先队列中
     */
    boolean offer(E e);

    /**
     * 返回优先队列优先级最高的元素,如果队列不存在元素则直接返回空
     */
    E poll();

    /**
     * 查看优先队列堆顶元素,如果队列为空则返回空
     */
    E peek();
}

public class PriorityQueue<E extends Comparable> implements Queue<E> {

    private MinHeap<E> data;


    public PriorityQueue() {
        data = new MinHeap<>();
    }

    public PriorityQueue(Comparator comparator) {
        data = new MinHeap<>(comparator);
    }


    @Override
    public int size() {
        return data.size();
    }

    @Override
    public boolean isEmpty() {
        return data.isEmpty();
    }

    @Override
    public boolean offer(E e) {
        data.add(e);
        return true;
    }

    @Override
    public E poll() {
        return data.poll();
    }

    @Override
    public E peek() {
        return data.peek();
    }
}

//往队列中随机添加1000_0000条数据
PriorityQueue<Integer> priorityQueue = new PriorityQueue<>();
int n = 1000_0000;
for (int i = 0; i < n; i++) {
    priorityQueue.offer(ThreadLocalRandom.current().nextInt(1, n));
}

//将优先队列中的数据按照优先级取出并存到数组中
int[] arr = new int[n];

for (int i = 0; i < n; i++) {
    arr[i] = priorityQueue.poll();
}

//如果前一个元素大于后一个元素,则说明优先队列优先级排列有问题
for (int i = 1; i < arr.length; i++) {
    if (arr[i - 1] > arr[i]) {
        throw new RuntimeException("error PriorityQueue");
    }
}

JDK 自带 PriorityQueue 使用示例

基本类型优先队列

import java.util.PriorityQueue;
public class Main {


    public static void main(String[] args) {
        //往队列中随机添加1000_0000条数据
        PriorityQueue<Integer> priorityQueue = new PriorityQueue<>();
        int n = 1000_0000;
        for (int i = 0; i < n; i++) {
            priorityQueue.offer(RandomUtil.randomInt(1, n));
        }

        //将优先队列中的数据按照优先级取出并存到数组中
        int[] arr = new int[n];

        for (int i = 0; i < n; i++) {
            arr[i] = priorityQueue.poll();
        }

        //如果前一个元素大于后一个元素,则说明优先队列优先级排列有问题
        for (int i = 1; i < arr.length; i++) {
            if (arr[i - 1] > arr[i]) {
                throw new RuntimeException("error PriorityQueue");
            }
        }

    }
}

特殊类型优先队列

public class Example {
    static class Person {
        String name;
        int age;

        Person(String name, int age) {
            this.name = name;
            this.age = age;
        }

        public String getName() {
            return name;
        }

        public int getAge() {
            return age;
        }
    }

    public static void main(String[] args) {
        PriorityQueue<Person> pq = new PriorityQueue<>(Comparator.comparing(Person::getAge));
        pq.add(new Person("Alice", 25));
        pq.add(new Person("Bob", 30));
        pq.add(new Person("Charlie", 20));

        while (!pq.isEmpty()) {
            Person p = pq.poll();
            System.out.println(p.getName() + " (" + p.getAge() + ")");
        }
    }
}

JDK 自带 PriorityQueue 源码分析

构造函数

public PriorityQueue(int initialCapacity,
                         Comparator<? super E> comparator) {
        //如果初始化容量小于1则抛出异常
        if (initialCapacity < 1)
            throw new IllegalArgumentException();
        this.queue = new Object[initialCapacity];
        this.comparator = comparator;
    }

public PriorityQueue(Collection<? extends E> c) {
		//SortedSet类型的集合转为PriorityQueue，从SortedSet中获取一个比较器进行初始化，然后按照比较器的规则转移元素到PriorityQueue中
        if (c instanceof SortedSet<?>) {
            SortedSet<? extends E> ss = (SortedSet<? extends E>) c;
            this.comparator = (Comparator<? super E>) ss.comparator();
            initElementsFromCollection(ss);
        }
        //如果集合类型是PriorityQueue则获取PriorityQueue的比较器，并将PriorityQueue的元素转移到我们的PriorityQueue中
        else if (c instanceof PriorityQueue<?>) {
            PriorityQueue<? extends E> pq = (PriorityQueue<? extends E>) c;
            this.comparator = (Comparator<? super E>) pq.comparator();
            initFromPriorityQueue(pq);
        }
        else {
            this.comparator = null;
            initFromCollection(c);
        }
    }

//SortedSet类型的集合转为PriorityQueue，从SortedSet中获取一个比较器进行初始化，然后按照比较器的规则转移元素到PriorityQueue中
        if (c instanceof SortedSet<?>) {
            SortedSet<? extends E> ss = (SortedSet<? extends E>) c;
            this.comparator = (Comparator<? super E>) ss.comparator();
            initElementsFromCollection(ss);
        }

private void initElementsFromCollection(Collection<? extends E> c) {
        Object[] a = c.toArray();
        // If c.toArray incorrectly doesn't return Object[], copy it.
        if (a.getClass() != Object[].class)
            a = Arrays.copyOf(a, a.length, Object[].class);
        int len = a.length;
        if (len == 1 || this.comparator != null)
            for (int i = 0; i < len; i++)
                if (a[i] == null)
                    throw new NullPointerException();
        this.queue = a;
        this.size = a.length;
    }

 //如果集合类型是PriorityQueue则获取PriorityQueue的比较器，并将PriorityQueue的元素转移到我们的PriorityQueue中
        else if (c instanceof PriorityQueue<?>) {
            PriorityQueue<? extends E> pq = (PriorityQueue<? extends E>) c;
            this.comparator = (Comparator<? super E>) pq.comparator();
            initFromPriorityQueue(pq);
        }

private void initFromPriorityQueue(PriorityQueue<? extends E> c) {
        if (c.getClass() == PriorityQueue.class) {
            this.queue = c.toArray();
            this.size = c.size();
        } else {
           //略
        }
    }

 else {
            this.comparator = null;
            initFromCollection(c);
        }

private void initFromPriorityQueue(PriorityQueue<? extends E> c) {
        if (c.getClass() == PriorityQueue.class) {
          //略
        } else {
            initFromCollection(c);
        }
    }

 private void initFromCollection(Collection<? extends E> c) {
        initElementsFromCollection(c);
        heapify();
    }

private void heapify() {
        for (int i = (size >>> 1) - 1; i >= 0; i--)
            siftDown(i, (E) queue[i]);
    }

入队操作

public boolean offer(E e) {
		//校验元素是否为空
        if (e == null)
            throw new NullPointerException();
        modCount++;
        //设置新插入的位置为size
        int i = size;
        //判断数组容量是否足够，如果不够则扩容
        if (i >= queue.length)
            grow(i + 1);
        size = i + 1;
        //如果是第一个元素，则直接将其放到索引0位置
        if (i == 0)
            queue[0] = e;
        else
        //如果不是第一个元素，则调用siftUp将元素放入队列
            siftUp(i, e);
        return true;
    }

private void siftUp(int k, E x) {
        if (comparator != null)
            siftUpUsingComparator(k, x);
        else
            siftUpComparable(k, x);
    }

private void siftUpUsingComparator(int k, E x) {
        while (k > 0) {
        //获取入队元素当前索引位置的父索引parent
            int parent = (k - 1) >>> 1;
            //根据父索引找到元素e
            Object e = queue[parent];
            // 如果新节点x比e大，则说明当前入队操作符合小顶堆要求，直接结束循环
            if (comparator.compare(x, (E) e) >= 0)
                break;
			//如果x比e小，则将父节点e的值改为我们入队元素x的值
            queue[k] = e;
            //k指向父索引，继续循环向上比较父索引，直到找到比x还小的父节点e，终止循环
            k = parent;
        }
        //将x存到符合要求的索引位置k
        queue[k] = x;
    }

出队操作

public E poll() {
        if (size == 0)
            return null;
        int s = --size;
        modCount++;
        E result = (E) queue[0];
        E x = (E) queue[s];
        queue[s] = null;
        if (s != 0)
            siftDown(0, x);
        return result;
    }

查看优先级最高的元素

 public E peek() {
        return (size == 0) ? null : (E) queue[0];
    }

Leetcode 中关于 PriorityQueue 的使用

问题

实现思路

public int[] topKFrequent(int[] nums, int k) {
        //统计各个元素出现的频率并存到map中
        HashMap<Integer, Integer> map = new HashMap<>();
        for (int num : nums) {
            if (!map.containsKey(num)) {
                map.put(num, 1);
            } else {
                map.put(num, map.get(num) + 1);
            }
        }

        //优先队列按照map的value进行升序排序
        PriorityQueue<Integer> queue = new PriorityQueue<>(Comparator.comparingInt(map::get));

        //遍历map,筛选出前k个高频元素
        map.entrySet().forEach(e -> {
            if (queue.size() < k) {
                queue.add(e.getKey());
            } else if (e.getValue() > map.get(queue.peek())) {
                queue.remove();
                queue.add(e.getKey());
            }
        });

        //转为数组返回
        int[] arr = queue.stream().mapToInt(Integer::intValue).toArray();
        return arr;

    }

public int[] topKFrequent(int[] nums, int k) {
        HashMap<Integer, Integer> map = new HashMap<>();
        for (int num : nums) {
            if (!map.containsKey(num)) {
                map.put(num, 1);
            } else {
                map.put(num, map.get(num) + 1);
            }
        }


        PriorityQueue<Integer> queue = new PriorityQueue<>(Comparator.comparingInt(map::get));

        map.entrySet().forEach(e -> {
            if (queue.size() < k) {
                queue.offer(e.getKey());
            } else if (e.getValue() > map.get(queue.peek())) {
                queue.poll();
                queue.offer(e.getKey());
            }
        });


        List<Integer> list = new ArrayList<>();
        while (!queue.isEmpty()) {
            list.add(queue.poll());
        }

        return list.stream().mapToInt(Integer::intValue).toArray();

    }

    public static void main(String[] args) {
        int[] nums = {1,1,32,312,3,432,412412412,44,4,444,444,4,4,4,4,4,4,4,4,4,6,6,6,6,6,4,4,4,44,3,4324,34,23432,4324,324,324,324,3242,34};
        int k = 2;
        for (int i : (new Solution()).topKFrequent(nums, k)) {
            System.out.println(i);
        }
    }

PriorityQueue 常见面试题

PriorityQueue 为什么要用二叉堆实现，二叉堆有什么使用优势吗？

PriorityQueue 是线程安全的吗？

PriorityQueue 底层是用什么实现的？初始化容量是多少？

private static final int DEFAULT_INITIAL_CAPACITY = 11;
public PriorityQueue() {
        this(DEFAULT_INITIAL_CAPACITY, null);
    }

如果我希望 PriorityQueue 按照从大到小的顺序排序要怎么做？

public static void main(String[] args) {
        //将比较器取反一下实现升序排列
        PriorityQueue<Integer> priorityQueue = new PriorityQueue<>(Comparator.comparingInt(Integer::intValue).reversed());
        //随机插入10个元素
        int n = 10;
        for (int i = 0; i < n; i++) {
            int num = RandomUtil.randomInt(1, n);
            System.out.println("add " + num);
            priorityQueue.offer(num);
        }

        //输出看看是否符合预期
        int s = priorityQueue.size();
        for (int i = 0; i < s; i++) {
            System.out.println(priorityQueue.poll());
        }

    }

为什么 PriorityQueue 底层用数组构成小顶堆而不是使用链表呢？

参考