PriorityQueue的数据结构

PriorityQueue是Java中的优先级队列，父类Queue（单端队列，即只能队列的一端进行插入删除元素）。PriorityQueue是根据元素优先级进行出入队，优先级由底层堆的实现来控制，默认小根堆，可以在PriorityQueue声明时自定义比较器来实现大根堆。

//建里大根堆，默认小根堆
PriorityQueue<Integer> heap = new PriorityQueue<>(new Comparator<Integer>() {
    @Override
    public int compare(Integer o1, Integer o2) {
        return o2 - o1;
    }
});

PriorityQueue的构造方法

使用PriorityQueue进行堆排序

public static void main(String[] args) throws CloneNotSupportedException {
    //建大根堆，默认小根堆
    PriorityQueue<Integer> heap = new PriorityQueue<>(new Comparator<Integer>() {
        @Override
        public int compare(Integer o1, Integer o2) {
            return o2 - o1;
        }
    });
    int[] arr = {6, 5, 4, 3, 2, 1};
    //将数组元素放入堆
    for (int i : arr) {
        heap.offer(i);
    }

    System.out.println(heap);
    int index = 0;
    //每次将堆顶元素放入数组，并删除堆顶元素，相当于每次从剩余堆中取最大值
    for (int i = 0; i < arr.length; i++) {
        arr[index++] = (Integer) heap.poll();
    }

    for (int i : arr) {
        System.out.print(i + " ");  //6 5 4 3 2 1
    }
}

使用PriorityQueue求解TopK

public int[] getLeastNumbers(int[] arr, int k) {
    if(arr == null || k == 0){
        return new int[0];
    }

    //优先级队列默认是小根堆。根据优先级出队
    PriorityQueue heap = new PriorityQueue();
    //arr元素入堆，并建队
    for(int i : arr){
        heap.offer(i);
    }

    int[] res = new int[k];
    int index = 0;
    for(int i = 0; i < res.length; i++){
        //每次从堆顶拿元素，并删除堆顶元素，堆自动调整
        res[index++] = (Integer)heap.poll();
    }
    return res;
}

补充：手写堆

public int[] getLeastNumbers(int[] arr, int k) {
    if(arr == null || k == 0){
        return new int[0];
    }

    int[] res = new int[k];
    int index = 0;
    heap_build(arr, arr.length);
    for(int i = arr.length - 1; i >= arr.length - k; i--){
        res[index++] = arr[0];
        swap(arr, 0, i);
        heapify(arr, i, 0);
    }

    return res;
}

public void heapify(int[] arr, int n, int i){
    if(i > n){
        return;
    }
    int c1 = 2 * i + 1;
    int c2 = 2 * i + 2;
    int min = i;

    if(c1 < n && arr[min] > arr[c1]){
        min = c1;
    }

    if(c2 < n && arr[min] > arr[c2]){
        min = c2;
    }

    if(min != i){
        swap(arr, min, i);
        heapify(arr, n, min);
    }
}

public void heap_build(int[] arr, int n){
    for(int i = (n - 1) / 2; i >= 0; i--){
        heapify(arr, n, i);
    }
}

public void swap(int[] arr, int i, int j){
    int temp = arr[i];
    arr[i] = arr[j];
    arr[j] = temp;
}