堆操作篇 -- 寻找第K大

题目描述

牛客网：NC88 寻找第K大
有一个整数数组，请你根据快速排序的思路，找出数组中第K大的数。

给定一个整数数组a ,同时给定它的大小n和要找的K(1<= K<=n) ，请返回第K大的数(包括重复的元素，不用去重)，保证答案存在。
要求时间复杂度 O(n)

题目解析

方法一：sort()排序

把数组元素排序，可直接得到第K大的值.

public int findKth(int[] a, int n, int K) {
    Arrays.sort(a);
    return a[n - K];      
}

时间复杂度：O(nlogn)
空间复杂度：O(1)

方法二：快排思想

使用快速排序思路，先找到初始值init，大的元素移到其左侧，小的元素移到其右侧，然后比较init当前位置：

• 当前下标为k - 1，说明找到了第K大的值，返回
• 当前下标小于k - 1，说明第K大的值在右侧子数组中，对右半段进行递归查找
• 当前下标大于k - 1，说明第K大的值在左侧子数组中，对左半段进行递归查找

快排本质上是分治思想，将问题分为一小段一小段进行解决。

public int findKth(int[] a, int n, int K) {
    // write code here
    return findK(a,0,n - 1,K);       
}
public int findK(int[] a,int left,int right,int k){
    if(left > right){
        return -1;
    }
    int pos = position(a,left,right);
    if(pos == k - 1){
        return a[pos];
    }else if(pos < k - 1){
        return findK(a,pos + 1,right,k);
    }else{
        return findK(a,left,pos - 1,k);
    }
}
public int position(int[] a,int left,int right){
    int init = a[left];
    while(left < right){
        while(left < right && a[right] <= init){
            right--;
        }
        a[left] = a[right];
        while(left < right && a[left] >= init){
            left++;
        }
        a[right] = a[left];
    }
    a[left] = init;
    return left;
}

时间复杂度：O(n)，用二分法缩短了时间，不用对所有元素进行排序
空间复杂度：O(n),递归栈最大深度

方法三：堆

求第K大问题可以用堆。构建最小堆，大小为K。最小堆堆顶元素最小，遍历数组，不断更新堆顶元素，最终堆顶元素就是第K大的元素。

public int findKth(int[] a, int n, int K) {
    // write code here
    PriorityQueue<Integer> pq = new PriorityQueue<>(K);
    for(int num : a){
        pq.offer(num);
        if(pq.size() > K){
            pq.poll();
        }                          
    }
    return pq.poll();
}

时间复杂度：O(nlogk)
空间复杂度：O(k),堆的开销