215. Kth Largest Element in an Array

Leetcode Divide and Conquer Heap

Find the $k$th largest element in an unsorted array. Note that it is the $k$th largest element in the sorted order, not the $k$th distinct element.

Example 1:

Input: [3,2,1,5,6,4] and k = 2
Output: 5

Example 2:

Input: [3,2,3,1,2,4,5,5,6] and k = 4
Output: 4

Note: You may assume $k$ is always valid, $1 ≤ k ≤$ array's length.

分析

这道题最简单的思路莫过于把输⼊的$n$个整数排序，排序之后位于最后⾯的第$k$个数就是最大的第$k$个数。时间复杂度是$n\log(n)$，空间复杂度为$O(1)$。

public int findKthLargest(int[] nums, int k) {
    if (nums == null || nums.length == 0 || k > nums.length)
        throw new IllegalArgumentException();
    Arrays.sort(nums);
    return nums[nums.length - k];
}

另一个方法是利用优先级队列保存前$K$个最大的数。时间复杂度是$O(n\log k)$，空间复杂度是$O(k)$. 该方法适用于海量数据，参考LeetCode 703. Kth Largest Element in a Stream。

public int findKthLargest(int[] nums, int k) {
    if (nums == null || nums.length == 0 || k > nums.length)
                throw new IllegalArgumentException();
    final PriorityQueue<Integer> pq = new PriorityQueue<>();
    for (int num : nums) {
        pq.offer(num);
        if (pq.size() > k) pq.poll();
    }
    return pq.peek();
}

比较好的方法是利用快速选择算法，专门用于选择最大/最小的第$k$个数。

public int findKthLargest(int[] nums, int k) {
    if (nums == null || nums.length == 0 || k > nums.length)
        throw new IllegalArgumentException();
    int lo = 0, hi = nums.length - 1;
    // 要调整K，不然求的是最小第K个数
    k = nums.length - k;
    while (lo < hi) {
        int j = partition(nums, lo, hi);
        if (j < k)          lo = j + 1;
        else if (j > k)     hi = j - 1;
        else                return nums[j];
    }
    return nums[k];
}

private int partition(int[] nums, int lo, int hi) {
    int pivot = nums[lo];
    int i = lo + 1, j = hi;
    while (true) {
        while (i < hi && nums[i] <= pivot) i++;
        while (j > lo && nums[j] >= pivot) j--;
        if (i >= j) break;
        swap(nums, i, j);
    }
    swap(nums, lo, j);
    return j;
}

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