TopK问题
说明:
1)求最小k个数,使用大顶堆;
2)求最大k个数,使用小顶堆。
实现代码如下
import java.util.ArrayList;
//第k小,则构造大顶堆
public class Main{
public static void main(String args[]) {
Main main=new Main();
int[] nums={1,2,3,4};
ArrayList<Integer> res=main.GetLeastNumbers_Solution(nums,4);
res.forEach(i-> System.out.println(i));
}
public ArrayList<Integer> GetLeastNumbers_Solution(int [] input, int k) {
if(k<=0 || input==null || input.length<k){
return new ArrayList<>();
}
int[] heap=new int[k];
for(int i=0;i<k;i++){
insert(heap,i,input[i]);
}
for(int i=k;i<input.length;i++){
if(heap[0]>input[i]){
heap[0]=input[i];
delete(heap,0,k);
}
}
ArrayList<Integer> res=new ArrayList<>();
for(int i=k-1;i>-1;i--)
res.add(heap[i]);
return res;
}
//插入
private void insert(int[] heap,int idx,int value){
heap[idx]=value;
while(idx!=0){
int parent=(idx-1)/2;
if(heap[parent]<heap[idx]){
swap(heap,parent,idx);
idx=parent;
}else{
break;
}
}
}
//删除
private void delete(int[] heap,int idx,int heapSize){
int left=idx*2+1;
int right=idx*2+2;
int largest=idx;
while(left<heapSize){
if(heap[left]>heap[idx])
largest=left;
if(right<heapSize && heap[right]>heap[largest])
largest=right;
if(largest!=idx){
//交换idx和largest索引处的值
swap(heap,largest,idx);
idx=largest;
left=2*idx+1;
right=2*idx+2;
}else{
break;
}
}
}
//交换
private void swap(int[] heap,int idx1,int idx2){
int tmp=heap[idx1];
heap[idx1]=heap[idx2];
heap[idx2]=tmp;
}
}
参考书籍:《程序员代码面试指南》-左程云