TopK问题算是面试中常考的,而且有实际价值的算法中比较有代表性的一个了,主要解决方法有堆、快速选择、排序三种方案。下一次专门讲一下快速选择类的算法。
解题思路有多种,比如可以对数组排序,使用Arrays.sort();然后取前面K个数,就是最小或者最大的K个数,sort可以指定按升序还是降序排列。当然,这么做的时间复杂度是Nlog(N),虽然的确很快,但是还是会被面试官拿来鞭尸。
其实他们想考的就是堆罢了,虽然都是Nlog(N)。使用堆可以得到topK问题的通用解法。就拿本题来讲,以最大的topK个数为例,
1、我们可以建一个小顶堆,当堆元素数量小于K个时,直接入堆,
2、当大于等于K个时,与堆顶部比较,如果大于堆顶,则堆顶出堆,然后执行入堆操作
3、将堆里面的元素全部弹出,即为最大的K个数
我们知道,堆一般都是用优先队列实现的,在Java里面,PriorityQueue默认就是小顶堆。
public int[] getTopK(int[] nums,int k){
PriorityQueue<Integer> heap=new PriorityQueue<>();
for(int i:nums){
if(heap.size()<k){
heap.offer(i);
}else if(heap.top()<i){
heap.poll();
heap.offer(i);
}
}
int [] result=new int[heap.size()];
for(int i=0;i<result.length;i++){
result[i]=heap.poll();
}
return result;
}
这个问题和上个问题不同之处就在于本题只需要一个数,直接套上题代码,弹出的第一个数就是第K大的数。
不过这时候面试官大概率会要求你写出一个复杂度为O(N)的算法,而上题复杂度Nlog(N),空间复杂度为k,就不好用了。
这时候考察的其实是一个快排的思想,快排还有个兄弟,叫快速选择,也就是quickSelect算法,时间复杂度可以降到O(N)。
还记得快排怎么做的吗?选取一个枢纽元,在一轮过后,左边元素全部小于等于枢纽元,右边元素全部大于等于枢纽元。然而我们找的是最小或者最大的第K个数,完全不用care另一部分的值,我们只对枢纽元左半部分或者右半部分递归的实现快速选择就行了。
经过快速选择后,所需要的数据就是数组中第k-1个元素。
//不考虑k不合法的情况了
public int getKthNumber(int nums[],int k){
quickSelect(nums,0,nums.length-1,k);
return nums[k-1];
}
private void quickSelect(int []nums,int left,int right,int k){
if(left<right){
int i=left,j=right;
int privot=nums[left];//选取枢纽元
while(i<j){
while(i<j&&nums[j]>=privot){j--;}//找到第一个小于枢纽元的值
while(i<j&&nums[i]<=privot){i++;}//找到第一个大于枢纽元的值
if(i<j){
swap(nums[i],i,j);
}
}
swap(nums,left,i);//恢复数纽元
if(k<=i){
quickSelect(nums,left,i-1,k);
}else{
quickSelect(nums,i+1,right,k);
}
}
}
private void swap(int[] nums,int i,int j){
int temp=nums[i];
nums[i]=nums[j];
nums[j]=temp;
}
这个问题用堆实现也比较简单,如果数组大小为奇数,我们设置一个(数组长度+1)/2的堆,然后按照题目1来弹出堆顶就是中位数。如果是偶数,我们设置一个大小为 数组长度/2 +1的一个堆,然后弹出两个元素取平均值即可。
然而这样时间复杂度是Nlog(N),面试官无法忍受。。。
想到题目2中快速选择的思想,我们可以这样做:
1、数组长度len为奇数,我们使用2中代码选择第(len+1)/2小的数,即为中位数
2、数组长度len为偶数,我们选择第len/2小的数和第len/2+1小的数,取均值即为中位数
话不多说,直接放码:
public void test(){
int[] nums={3,5,1,6,7};
if(nums.length%2==0){
int k1=(nums.length+1)>>1;
quickSelect(nums,k1);
int a=nums[k1-1];
int k2=(nums.length+1)>>1+1;
quickSelect(nums,k2);
int b=nums[k2-1];
System.out.println((a+b)>>1);
}else {
int k=(nums.length+1)>>1;
quickSelect(nums,k);
System.out.println(nums[k-1]);
}
}
public void quickSelect(int[] nums,int k){
quickSelect(nums,0,nums.length-1,k);
}
public void quickSelect(int nums[],int left,int right,int k){
if(left<right){
int i=left,j=right;
int pivot=nums[left];
while(i<j){
while(i<j&&nums[j]>=pivot){j--;}
while(i<j&&nums[i]<=pivot){i++;}
if(i<j){
swap(nums,i,j);
}
}
swap(nums,left,i);
if(k<=i){
quickSelect(nums,left,i-1,k);
}else {
quickSelect(nums,i+1,right,k);
}
}
}
private void swap(int[] nums, int i, int j) {
int temp=nums[i];
nums[i]=nums[j];
nums[j]=temp;
}