java笔试题：海量数据找最大或最小的k个数（堆排序）

题目

海量数据找最大或最小的k个数，这里以找最小的K个数为例

堆排序

例如给一个数组nums【】这棵树就是完全二叉树，则：

• nums【i】的左节点为：num【2 * i + 1】
• nums【i】的右节点为：num【2 * i + 2】

这里简单说一下堆排序的流程：

• 第一步：将该数组调整为一个大（小）顶堆。即根节点元素永远大于左右节点。调整顺序是从倒数第一个非叶子节点（index = num.length / 2 - 1）开始。
  • 进入循环，index = num.length / 2 - 1
  • 调整堆
  • index–
  • 出循环
• 第二步：每次将堆顶元素和堆尾元素交换，数组长度减一，重新调整。每次都将最大的元素放到了最后。
  • 进入循环。
  • 将堆顶元素与堆尾元素交换。数组长度减一。因为已经将最大的元素定位到堆尾了。
  • 从堆顶重新调整堆为大顶堆。
  • 出循环

代码如下：

public class test16 {

    public static void main(String[] args) {
        int[] arr = new int[]{4,6,8,5,9,50};
        heapSort(arr);
        System.out.println(Arrays.toString(arr));
    }


    public static void heapSort(int[] arr){
        if(arr == null || arr.length <= 2) return;

        //创建初始大顶堆  根节点的元素一定大于子节点
        for(int i = arr.length/2 - 1; i >= 0; i--){
            adjustHeap(arr,i,arr.length);
        }
        
        
        //每循环一次，都会将队尾的元素定位 堆的长度就减1
        for(int i = arr.length - 1; i >= 0; i--){
            swap(arr,0,i); //交换堆顶 和堆尾元素
            adjustHeap(arr,0,i);//重新调整堆  但是此时堆的长度是i 不是i+1
        }
    }


    public static void adjustHeap(int[] arr, int index, int length){

        for(int left = 2 * index + 1; left < length; left = 2 * index + 1){
            int right = left + 1; //右子节点
            if(right < length && arr[right] > arr[left]){//比较左右子节点哪个比较大  就换哪个
                left = right;
            }
            //此时left指向当前节点左右子节点较大的那个节点
            if(arr[left] > arr[index]){
                swap(arr,left,index);
                index = left; //这次变换可能会导致该子节点不满足大小顶堆的要求，因此index要转向该子节点继续调整
            }else {
                break;
            }

        }

    }

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

}

题解

思路：

首先维护一个k个元素的大顶堆，然后遍历arr中的后续元素，比较当前元素是否小于堆顶元素，如果小于就进行替换，然后重新调整大顶堆。最终这个大顶堆就是最小的k个元素。

实现方式：

• 可以直接用数组实现，自己写堆排序，堆调整
• java中有已经包装好的堆结构，直接使用。 Java 中提供了现成的 PriorityQueue（默认小根堆），但是可以设置参数为大顶堆。

代码如下：

数组实现：

class Solution {
    public int[] getLeastNumbers(int[] arr, int k) {
        if(k == 0) return new int[0];
        int[] res = new int[k];

        //定义初始堆k个元素
        for(int i = 0; i < k; i++){
            res[i] = arr[i];
        }

        //初始化大顶堆
        for(int i = res.length / 2 - 1; i >= 0; i--){  
            adjustSort(res,i,res.length);
        }

        //每次比较当前的堆顶和arr中的元素
        for(int i = k; i < arr.length; i++){
            if(arr[i] < res[0]) { //如果当前元素小于堆顶，那就替换，再调整堆
                res[0] = arr[i];
                adjustSort(res,0,res.length);
            }
        }

        return res;
        
    }

    //调整堆
    public void adjustSort(int[] nums, int index, int length){

        for(int left = 2 * index + 1; left < length; left = 2 * index + 1){
            int right = left + 1;
            //比较当前节点和左右子节点哪个大 left就指向哪个 
            if(right < length && nums[right] > nums[left]){
                left = right;
            }

            //如果当前left指向的节点大于当前节点 就替换
            //同时替换之后会导致替换的子节点可能不满足大顶堆条件，因此index转向该子节点，继续调整
            if(nums[left] > nums[index]){
                swap(nums,left,index);
                index = left; 
            }else break;

        }
    }

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

java PriorityQueue 实现

class Solution {
    public int[] getLeastNumbers(int[] arr, int k) {
        if (k == 0 || arr.length == 0) {
            return new int[0];
        }
        // 默认是小根堆，实现大根堆需要重写一下比较器。
        Queue<Integer> pq = new PriorityQueue<>((v1, v2) -> v2 - v1);
        for (int num: arr) {
            if (pq.size() < k) { //先填k个元素进去
                pq.offer(num);
            } else if (num < pq.peek()) { //如果小于堆顶 
                pq.poll();//加入堆
                pq.offer(num);//将堆顶元素出了
            }
        }
        
        // 返回堆中的元素
        int[] res = new int[pq.size()];
        int idx = 0;
        for(int num: pq) {
            res[idx++] = num;
        }
        return res;
    }
}