堆排序与取topK java实现

1.堆排序思路

最近趁着有点时间，稍微复习了一下数据结构相关内容，温习了一下堆排序，做一下记录。

首先我们复习一下什么是堆：
堆是具有以下性质的完全二叉树：每个结点的值都大于或等于其左右孩子结点的值，称为大顶堆；或者每个结点的值都小于或等于其左右孩子结点的值，称为小顶堆。

堆排序的基本思想如下：
将待排序序列构造成一个大顶堆，此时，整个序列的最大值就是堆顶的根节点。将其与末尾元素进行交换，此时末尾就为最大值。然后将剩余n-1个元素重新构造成一个堆，这样会得到n个元素的次小值。如此反复执行，便能得到一个有序序列了。

更具体的过程与图示，见参考文献1，不再重新画图。

2.java实现

下面我们用java来实现一下堆排序的过程。

public class HeapSort {

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

    public static void adjust(int[] arr, int start, int end) {
        // 左右子节点
        int left = start * 2 + 1;
        int right = start * 2 + 2;
        int largest = start;
        if (left <= end && arr[left] > arr[largest]) largest = left;
        if (right <= end && arr[right] > arr[largest]) largest = right;

        // 交换数据，并继续调整
        if (largest!= start) {
            swap(arr, start, largest);
            adjust(arr, largest, end);
        }
    }

    // 从第一个非叶子结点开始调整，注意顺序是从下往上
    public static void buildHeap(int[] arr) {
        for(int i = arr.length / 2 - 1; i >= 0; i--) {
            adjust(arr, i, arr.length - 1);
        }
    }

    // 先构建堆，此时堆顶为最大元素，交换到此时数组的最后一位。
    public static void heapSort(int[] arr) {
        buildHeap(arr);
        for(int i = arr.length - 1; i > 0; i--) {
            swap(arr, 0, i);
            adjust(arr, 0, i - 1);
        }
    }

    public static void main(String[] args) {
        int[] arr = {1, 3, 5, 6, 2, 4, 6, 9, 8, 10, 7};
        heapSort(arr);
        for(int i = 0; i < arr.length; i++) {
            System.out.print(arr[i] + " ");
        }
    }
}

关键的步骤以及作用，都已经在代码中进行了注释，再结合参考文献1就可以容易理解。

3.求topK

求一个无序序列的topK，是个经典问题。这个经典问题的经典解法，就包括堆排序。

    public static void findTopK(int[] arr, int k) {
        PriorityQueue<Integer> pq = new PriorityQueue();
        for (int i = 0; i < k; i++) {
            pq.offer(arr[i]);
        }
        for (int i = k; i < arr.length; i++) {
            int current = pq.peek();
            if (arr[i] > current) {
                pq.poll();
                pq.offer(arr[i]);
            }
        }

        Integer[] ret = pq.toArray(new Integer[k]);
        Arrays.sort(ret, Collections.reverseOrder());

        for(int i = 0; i < k; i++) System.out.print(ret[i] + " ");
    }

    public static void findLastK(int[] arr, int k) {
        PriorityQueue<Integer> pq = new PriorityQueue<Integer>((Integer o1, Integer o2) -> o2 - o1);
        for (int i = 0; i < k; i++) {
            pq.offer(arr[i]);
        }
        for (int i = k; i < arr.length; i++) {
            int current = pq.peek();
            if (arr[i] < current) {
                pq.poll();
                pq.offer(arr[i]);
            }
        }
        Integer[] ret = pq.toArray(new Integer[k]);
        Arrays.sort(ret);
        for(int i = 0; i < k; i++) System.out.print(ret[i] + " ");
    }

    public static void main(String[] args) {
        int[] arr = {1, 3, 5, 10, 8, 6, 7, 9, 2, 4};
        int k = 3;
        findTopK(arr, k);
        System.out.println();
        findLastK(arr, k);
    }

上面的代码，分别找到最大的三个数与最小的三个数。
PriorityQueue 是java中的优先队列，默认就是小顶堆实现。求top3最大值时候，用小顶堆即可。如果求top3最小值，则使用大顶堆。

上面代码运行最后的输出：