快速选择算法

LeetCode215数组中的第K个最大元素

题目描述

给定整数数组 nums 和整数 k，请返回数组中第 k 个最大的元素。

请注意，你需要找的是数组排序后的第 k 个最大的元素，而不是第 k 个不同的元素。

你必须设计并实现时间复杂度为 O(n) 的算法解决此问题。

示例 1:

输入: [3,2,1,5,6,4], k = 2
输出: 5

示例 2:

输入: [3,2,3,1,2,4,5,5,6], k = 4
输出: 4

提示：

• 1 <= k <= nums.length <= 105
• -104 <= nums[i] <= 104

方法一（小顶堆）

class Solution {
    public int findKthLargest(int[] nums, int k) {
        PriorityQueue<Integer> pq = new PriorityQueue<>(); // min heap
        for (int val : nums) {
            pq.add(val);
            if (pq.size() > k) {
                pq.poll();
            }
        }
        return pq.peek();
    }
}

在这个代码中，我们首先创建一个优先队列pq。然后，我们遍历数组nums，对于每个元素，我们将其添加到队列中。如果队列的大小大于k，我们就删除队列中的最小元素。最后，我们返回队列的头部元素，这就是第k个最大的元素。

这个算法的时间复杂度是O(n log k)，因为我们需要对每个元素进行堆操作，堆操作的时间复杂度是O(log k)。虽然题目要求时间复杂度为O(n)，但是这个算法在实际中的性能已经非常好了，因为log k通常远小于n。

方法二（快速选择）

class Solution {
    public void swap(int[] nums,int a,int b){
        int tem = nums[a];
        nums[a] = nums[b];
        nums[b] = tem;
    }
    public int getIndex(int[] nums,int left,int right){

        int m = (right+left) / 2 ;
        int n = nums[right];
        int p = left;
        for(int i = left;i<right;i++){
            if(nums[i] < n){
                swap(nums,p,i);
                p++;
            }
        }
        swap(nums,p,right);
        return p;
    }
    public int quickSelect(int[] nums,int left,int right,int k){
        if(right == left){
            return nums[left];
        }
        int index = getIndex(nums,left,right);
        if(index == k){
            return nums[index];
        }else if(index < k){
            return quickSelect(nums,index+1,right,k);
        }else{
            return quickSelect(nums,left,index-1,k);
        }
    }
    public int findKthLargest(int[] nums, int k) {
        return quickSelect(nums,0,nums.length-1,nums.length - k);
    }
}

快速选择是快速排序的一个优化，它在平均情况下的时间复杂度是O(n)。

在这个代码中，我们首先使用快速选择算法来找到第k个最大的元素。快速选择算法的基本思想是，我们选择一个枢轴元素，然后将数组分为两部分，一部分是小于枢轴的元素，另一部分是大于枢轴的元素。然后我们根据k和枢轴的位置，决定是在左边的部分还是右边的部分继续查找。如果k等于枢轴的位置，那么枢轴就是我们要找的元素。