leetcode 215. 数组中的第K个最大元素

题目描述

在未排序的数组中找到第 k 个最大的元素。请注意，你需要找的是数组排序后的第 k 个最大的元素，而不是第 k 个不同的元素。
相关标签： 堆、分治算法   难度： 中等

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

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

说明:
你可以假设 k 总是有效的，且 1 ≤ k ≤ 数组的长度。

思路：
见bfprt算法。

class Solution {
    public int findKthLargest(int[] nums, int k) {
        return bfprt(nums, 0, nums.length - 1, k - 1);
    }
    //排序，然后返回中位数
    private int getMedian(int[] nums, int begin, int end){
        for(int i = begin + 1;i <= end;i++){
            for(int j = i;j > begin && nums[j] < nums[j - 1]; j--){
                swap(nums, j ,j - 1);
            }
        }
        return nums[begin + ((end - begin) >> 1)];
    }
    private int getMedianOfMedians(int[] nums, int begin, int end){
        int num = end - begin + 1;
        int offset = num % 5 == 0 ? 0 : 1;
        int[] medians = new int[num / 5 + offset];
        for(int i = 0; i < medians.length;i++){
            int subBegin = begin + i * 5;
            int subEnd = subBegin + 4;
            medians[i] = getMedian(nums, subBegin, Math.min(subEnd, end));
        }
        //调用bfprt方法找到medians的中位数
        return bfprt(medians, 0, medians.length - 1, medians.length / 2);
    } 
    private int bfprt(int[] nums, int begin, int end, int k){
        if(begin == end) return nums[begin];
        int pivot = getMedianOfMedians(nums, begin, end);
        int[] partition = partition(nums, begin, end, pivot);
        if(k >= partition[0] && k <= partition[1]){
            return nums[k];
        }else if(k > partition[1]){
            return bfprt(nums, partition[1] + 1, end, k);
        }else{
            return bfprt(nums, begin, partition[0] - 1, k);
        }
    }
    
    private int[] partition(int[] nums, int begin, int end, int pivot){
        int l = begin - 1, r = end + 1;
        int i = begin;
        while(i < r){
            if(nums[i] == pivot){
                i++;
            }else if(nums[i] > pivot){
                swap(nums, ++l, i++);
            }else{
                swap(nums, i, --r);
            }
        }
        return new int[]{l + 1, r - 1};
    }
    private void swap(int[] nums, int i, int j){
        int temp = nums[i];
        nums[i] = nums[j];
        nums[j] = temp;
    }
}