一、题目内容

image.png

二、题解

解法1：小顶堆

思路

堆。
PriorityQueue 默认为小顶堆，实现Comparator可以改为大顶堆。

代码

  1. int[] ans = new int[k];
  2.         if (k == 0) {
  3.             return ans;
  4.         }
  6.         PriorityQueue<Integer> queue = new PriorityQueue<Integer>(new Comparator<Integer>() {
  7.             @Override
  8.             public int compare(Integer o1, Integer o2) {
  9.                 //大顶
  10.                 return o2 - o1;
  11.             }
  12.         });
  14.         //先填满
  15.         for (int i = 0; i < k; i++) {
  16.             queue.add(arr[i]);
  17.         }
  19.         for (int i = k; i < arr.length; i++) {
  20.             if (queue.isEmpty()) {
  21.                 queue.add(arr[i]);
  22.             } else {
  23.                 if (queue.peek() > arr[i]) {
  24.                     queue.poll();
  25.                     queue.add(arr[i]);
  26.                 }
  27.             }
  28.         }
  30.         for (int i = 0; i < k; i++) {
  31.             ans[i] = queue.poll();
  32.         }
  34.         return ans;

解法2：大顶堆，自己实现

思路

用大顶堆，自己实现一个堆

代码

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

    buildMaxHeap(maxHeapArr);
    while (i < arr.length) {
        if (arr[i] < maxHeapArr[0]) {
            maxHeapArr[0] = arr[i];
            adjustMaxHeap(maxHeapArr, 0);
        }
        i++;
    }

    return maxHeapArr;
}

public void buildMaxHeap(int[] a) {
    for (int i = a.length / 2 - 1; i >= 0; i--) {
        adjustMaxHeap(a, i);
    }
}

public void adjustMaxHeap(int[] a, int currPos) {
    int len = a.length;

    int tempRoot = a[currPos];
    //从当前节点的左子节点开始
    for (int k = currPos * 2 + 1; k < len; k = 2 * k + 1) {
        //如果左子节点小于右子节点，则右移一位，然后比较最大的子节点与当前的跟大小，小于当前跟则替换，使其满足大顶堆条件
        if (k + 1 < len && a[k] < a[k + 1]) {
            k++;
        }
        if (a[k] > tempRoot) {
            a[currPos] = a[k];
            currPos = k;
        } else {
            break;
        }
    }
    a[currPos] = tempRoot;
}