解法一:使用PriorityQueue

import java.util.*;
public class Solution {
    public ArrayList<Integer> GetLeastNumbers_Solution(int [] input, int k) {
        if(k>input.length||k<=0) return new ArrayList<>();
        PriorityQueue<Integer> pq=new PriorityQueue<>((o1,o2)->o2-o1);
        for(int i:input){
            if(pq.size()==k&&pq.peek()>=i) pq.poll();
            else if(pq.size()==k&&pq.peek()<i) continue;
            pq.add(i);
        }
        ArrayList<Integer> list=new ArrayList<>(pq.size());
        while(!pq.isEmpty()) list.add(pq.poll());
        Collections.reverse(list);
        return list;
    }
}

解法二:大根堆

实质上,是自己手写了一遍PriorityQueue的内容,堆排序就是它的原理。

import java.util.*;
public class Solution {
    public ArrayList<Integer> GetLeastNumbers_Solution(int [] input, int k) {
        if(input==null||k<=0||k>input.length) return new ArrayList<Integer>();
        //划定input前k个元素为备选元素
        //我们会把可能为答案的值不断地交换到这里
        createBigHeap(input, k);
        //如果发现有元素比备选元素的最大值input[0]要小
        //说明备选元素最大值一定不会是最小的K个数之一,所以把它们交换一下
        for(int i=k; i<input.length; i++){
            if(input[0]>input[i]){
                swap(input, 0, i);
                ajustDown(input, 0, k);
            }
        }    
        ArrayList<Integer> ans=new ArrayList<>(k);
        for(int i=0; i<k&&i<input.length; i++)
            ans.add(input[i]);
        return ans;
    }

    //交换数组input位置i和j上的元素
    void swap(int[] input, int i, int j){
        int t=input[i];
        input[i]=input[j];
        input[j]=t;
    }

    //构造大根堆,即,对于每一个非叶子节点input[i]
    //一定大于它的左孩子input[2*i+1]和右孩子input[2*i+2]
    //如此,根节点一定会获得最大值,所以叫“大”根堆
    void createBigHeap(int[] input, int k){
        for(int i=k/2; i>=0; i--)
            ajustDown(input, i,k);
        return;
    }

    //把位置i上的元素放到它在大根堆中的合理位置
    //即它的父节点都比它大,子节点都比它小
    //由于插入时都是直接从根插入的,所以只需要往叶子节点交换即可
    //这一步通过比较input[i], input[left_child], input[right_child]实现
    void ajustDown(int[] input, int i, int k){
        if(2*i+1>=k) return;
        int left=2*i+1, right=2*i+2;
        if(input[i]<input[left]){
            if(right<k&&input[left]<input[right]){
                swap(input, i, right);
                ajustDown(input, right, k);
            }else{
                swap(input, i, left);
                ajustDown(input, left, k);
            }
        }else if(right<k&&input[i]<input[right]){
            swap(input, i, right);
            ajustDown(input, right, k);
        }
        return;
    }
}

FAQ.

Q: 大根堆是什么?
A: 它的实质是一棵满二叉树(不一定完全)。其中父节点一定大于它的所有子节点。对于这样的一棵满二叉树,它的层序遍历是“连续的”。并且若父节点为层序遍历的第i位,左孩子和右孩子一定分别为 2i+1, 2i+2。因此,我们可以用一维数组来表示这个大根堆。注意,大根堆只对父子节点之间的关系有要求,对左右孩子的大小关系没有任何要求。因此在构建大根堆的时候,只用从父节点依次向上检查即可。(这就是为什么createBigHeap中循环是从 i=k/2 开始的缘故。)

Q: ajustDown是怎样实现的?
A: 它的作用是把input[i]元素放到大根堆适当的地方去。因为元素都是从根节点插入的,所以只用把input[i]和它的孩子交换就行了。它的孩子分别为在位置left=2i+1, right=2i+2的元素。这里写了一些if else都是为了比较它和左右孩子的大小关系。跟左右孩子对换之后,可能还要进一步往叶子节点对换,所以这里递归了。

Q: 如何记忆大根堆的写法?
A:

  1. 尽可能地把不同功能的代码块分开, createBigHeap, ajustDown, swap写成三个函数,左右孩子明明白白写成left,right,别写成2i+1, 2i+2看着头晕。
  2. 范围范围范围!createBigHeap只检查非叶子节点,每次查找左右孩子的值前,请确认left,right都在数组范围内。
  3. 大根堆用来找最小的K个数,小根堆用来找最大的K个数。