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

题目描述

在未排序的数组中找到第 k 个最大的元素。

示例

输入：[3,2,1,5,6,4]，k = 2

输出：5

解法一：排序

思路：

对原数组进行排序，然后输出排序后第 k 个数即可。

时间复杂度为 $O(n\log n)$，其中 n 是数组的长度。

代码：（排序，时间复杂度 $O(n\log n)$）

class Solution {
    public int findKthLargest(int[] nums, int k) {
        Arrays.sort(nums);
        return nums[nums.length - k];
    }
}

解法二：堆

思路：

建立一个大小为 k 的小根堆，遍历数组，将每个元素与小根堆的堆顶元素比较，如果大于堆顶元素，则将堆顶元素替换成该元素，并重新调整小根堆，保证堆的大小仍为 k。

时间复杂度为 $O(n\log k)$，其中 n 是数组的长度。

代码：（堆，时间复杂度 $O(n\log k)$）

class Solution {
    public int findKthLargest(int[] nums, int k) {
        PriorityQueue<Integer> pq = new PriorityQueue<>(k);
        for (int num : nums) {
            if (pq.size() < k) {
                pq.offer(num);
            } else if (num > pq.peek()) {
                pq.poll();
                pq.offer(num);
            }
        }
        return pq.peek();
    }
}

解法三：快速选择

思路：

使用快速排序中的 partition 函数，将数组分为左右两部分，左边的元素都大于等于基准值，右边的元素都小于等于基准值。比较基准值的下标 left 和 k 的大小，如果 left 等于 k，说明基准值即为第 k 大的元素；如果 left 小于 k，说明第 k 大的元素在右边部分中，以此递归处理右边部分；如果 left 大于 k，说明第 k 大的元素在左边部分中，以此递归处理左边部分。

平均时间复杂度为 $O(n)$，最坏情况时间复杂度为 $O(n^2)$。

代码：（快速选择，平均时间复杂度 $O(n)$，最坏情况时间复杂度 $O(n^2)$）

class Solution {
    private int partition(int[] nums, int l, int r) {
        int pivot = nums[r];
        int i = l - 1;
        for (int j = l; j < r; ++j) {
            if (nums[j] >= pivot) {
                ++i;
                swap(nums, i, j);
            }
        }
        swap(nums, i + 1, r);
        return i + 1;
    }

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

    public int findKthLargest(int[] nums, int k) {
        int l = 0, r = nums.length - 1;
        while (l <= r) {
            int pos = partition(nums, l, r);
            if (pos == k - 1) {
                return nums[pos];
            } else if (pos < k - 1) {
                l = pos + 1;
            } else {
                r = pos - 1;
            }
        }
        return -1;
    }
}