牛客题霸NC88题解

寻找第K大

牛客题霸NC88

难度：Medium

题目描述

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

给定一个整数数组a,同时给定它的大小n和要找的K(K在1到n之间)，请返回第K大的数，保证答案存在。

示例1

输入

[1,3,5,2,2],5,3

返回值

2

题目解答

利用快排思想

通过快速排序的partion与二分思想找到第K大的数，代码如下：

import java.util.*;

public class Finder {
    public int findKth(int[] a, int n, int K) {
        // write code here
        
        int low = 0, high = n-1;
        K = n - K;
        
        int index = findKthCore(a, low, high, K);
        return a[index];
        
    }
    
    public int findKthCore(int[] a, int low, int high, int K){
        
        int mid = partion(a, low, high);
        
        if(mid - low == K){
            return mid;
        }
        if(mid - low > K){
            return findKthCore(a, low, mid-1, K);
        }
        else{
            return findKthCore(a, mid+1, high, K - (mid - low + 1));
        }
        
    }
    
    public int partion(int[] a, int low, int high){
        
        while(low < high){
            while(low < high && a[high] >= a[low]){
                high--;
            }
            swap(a, low, high);
            while(low < high && a[low] <= a[high]){
                low++;
            }
            swap(a, low, high);
        }
        
        return low;
        
    }
    
    public void swap(int[] a, int i, int j){
        
        int temp = a[i];
        a[i] = a[j];
        a[j] = temp;
    }
}

时间复杂度： O(nlogn)

空间复杂度：O(1)

通过堆实现

使用一个大小为K的小顶堆，将数组中的所有元素依次加入堆中，最终堆顶元素就是第K大的数。

import java.util.*;

public class Finder {
    public int findKth(int[] a, int n, int K) {
        // write code here
        
        // 小顶堆
        PriorityQueue<Integer> priorityQueue = new PriorityQueue<>();
        
        for(int m : a){
            // 堆元素数量不够K个或者待入堆元素比堆顶元素大
            if(priorityQueue.size() < K || priorityQueue.peek() <= m){
                priorityQueue.offer(m);
            }
            // 保持堆大小为K
            if(priorityQueue.size() > K){
                priorityQueue.poll();
            }
        }
        
        return priorityQueue.peek();
        
    }
}

时间复杂度： O(nlogK)

空间复杂度：O(K)