快速选择：数组中第k个最大元素

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

题目描述

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

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

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

示例

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

解题思路

快排思想快速选择第k大的元素。

• 根据pivot将数组分为三部分
• 根据下标k落在数组的哪一部分，只处理这一部分
  • 如果在 等于pivot的部分，不需要关注顺序直接返回。
  • 如果在 小于/大于pivot的部分，递归快速选择这一部分，因为这部分只是全部小于/大于pivot 但是整体是无序的。

完整代码：

class Solution {
    public int findKthLargest(int[] nums, int k) {
        return quickSelect(nums, 0, nums.length - 1, nums.length - k);
    }

    int quickSelect(int[] nums, int start, int end, int k) {
        Random random = new Random();
        int pivot = nums[start + random.nextInt(end - start + 1)];

        int i = start, l = start, r = end;
        while (i <= r) {
            if (nums[i] == pivot) {
                i++;
            }
            else if (nums[i] < pivot) {
                swap(nums, l++, i++);
            }
            else {
                swap(nums, i, r--);
            }
        }

        if (k >= l && k <= r) {
            return nums[k];
        }
        else if (k < l) {
            return quickSelect(nums, start, l - 1, k);
        }
        else {
            return quickSelect(nums, i, end, k);
        }
    }

    void swap(int[] nums, int l, int r) {
        int tmp = nums[l];
        nums[l] = nums[r];
        nums[r] = tmp;
    }
}