题目描述

给定整数数组 nums 和整数 k,请返回数组中第 k 个最大的元素。
请注意,你需要找的是数组排序后的第 k 个最大的元素,而不是第 k 个不同的元素。

示例1:

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

示例2:

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

思路

该题为快速排序partition的运用:

  • 找第K个最大元素,就是找排序完成后第nums.length - k个元素的位置;
    • 例如k = 3, 1, 2, 3, 4, 5, 6, 7, 要找的元素是5, 就是第 7 - 3 = 4个元素的位置 (targetIndex)。
  • 做一次partition, 可以找到最右边元素pivot的最终位置finalIndex,并且左边的元素比它小(或相等),右边的元素比它大;
    • 如果发现finalIndex == targetIndex, nums[targetIndex];
    • 如果finalIndex < targetIndex, 说明要找的元素比pivot大,那么就去finalIndex的右边找;
    • 否则,去finalIndex的左边找。
  • partition的具体实现,参考快速排序

代码

public class Solution_215 {
    public static void main(String[] args) {
        int[] nums = {99, 99};
        int k = 1;
        Solution_215 obj = new Solution_215();
        System.out.println(obj.findKthLargest(nums, k));

    }

    public int findKthLargest(int[] nums, int k) {
        int targetIndex = nums.length - k;

        int left = 0;
        int right = nums.length - 1;

        while(true) {
            int pivotIndex = partition(nums, left, right);
            if (targetIndex == pivotIndex) {
                break;
            } else if (targetIndex > pivotIndex) {
                left = pivotIndex + 1;
            } else {
                right = pivotIndex - 1;
            }
        }

        return nums[targetIndex];
    }

    private int partition(int[] nums, int left, int right) {
        int randomIndex = new Random().nextInt(right - left + 1) + left;
        swap(nums, randomIndex, right);
        int pivot = nums[right];

        int i = left - 1;

        for (int j = left; j < right; j++) {
            if (nums[j] <= pivot) {
                i++;
                swap(nums, i, j);
            }
        }

        swap(nums, i + 1, right);

        return i + 1;
    }

    private void swap(int[] nums, int i, int j) {
        int temp = nums[i];
        nums[i] = nums[j];
        nums[j] = temp;
    }
}