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