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

本题思路根据 wisdompeak 大神的 题解 整理。

题目

题目链接

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

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

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

示例 1:

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

示例 2:

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

提示：

• 1 <= k <= nums.length <= 10^5
• -10^4 <= nums[i] <= 10^4

方法 1：堆排序

思路

维护一个容量为 $k$ 的小顶堆。遍历数组，考虑以下两种情况：

1. 小顶堆的元素个数小于 $k$，直接加入当前元素。
2. 小顶堆的元素个数等于 $k$，如果当前元素大于堆顶元素，则移除堆顶元素并加入当前元素。

复杂度分析

• 时间复杂度: $\mathcal{O}(n\log k)$。
• 空间复杂度: $\mathcal{O}(k)$。

代码：直接使用优先队列

public class Solution215 {  
    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 (pq.peek() < num) {  
                pq.poll();  
                pq.offer(num);  
            }  
        }  
        return pq.peek();  
    }  
}

代码：手写堆排序

public class Solution215 {  
    public int findKthLargest(int[] nums, int k) {  
        buildMinHeap(nums, k);  
        for (int i = k; i < nums.length; i++) {  
            if (nums[i] < nums[0]) {  
                continue;  
            }  
            swap(nums, 0, i);  
            minHeapify(nums, 0, k);  
        }  
        return nums[0];  
    }  
  
    private void minHeapify(int[] nums, int root, int size) {  
        int left = root * 2 + 1, right = left + 1, min = root;  
        if (left < size && nums[left] < nums[min]) {  
            min = left;  
        }  
        if (right < size && nums[right] < nums[min]) {  
            min = right;  
        }  
        if (min != root) {  
            swap(nums, root, min);  
            minHeapify(nums, min, size);  
        }  
    }  
  
    private void swap(int[] nums, int i, int j) {  
        int temp = nums[i];  
        nums[i] = nums[j];  
        nums[j] = temp;  
    }  
  
    private void buildMinHeap(int[] nums, int k) {  
        for (int i = k / 2 - 1; i >= 0; i--) {  
            minHeapify(nums, i, k);  
        }  
    }  
}

方法 2：二分搜索

思路

使用二分搜索模板。

复杂度分析

• 时间复杂度: $\mathcal{O}(n\log C)$。在本题的数据范围下，$C$ 约为 15。
• 空间复杂度: $\mathcal{O}(1)$。

代码

public class Solution215 {
    public int findKthLargest(int[] nums, int k) {  
        int left = -10001, right = 10001;  
        while (left + 1 < right) {  
            int mid = (left + right) >> 1;  
            if (count(nums, mid) >= k) {  
                left = mid;  
            } else {  
                right = mid;  
            }  
        }  
        return left;  
    }  
  
    private int count(int[] nums, int target) {  
        int res = 0;  
        for (int num : nums) {  
            res += num >= target ? 1 : 0;  
        }  
        return res;  
    }  
}

方法 3：快速选择

思路

使用快速选择模板。

复杂度分析

• 时间复杂度: $\mathcal{O}(n)$。
• 空间复杂度: $\mathcal{O}(\log n)$。

代码

public class Solution215 {  
    public int findKthLargest(int[] nums, int k) {  
        return quickSelect(nums, k, 0, nums.length - 1);  
    }  
  
    private int quickSelect(int[] nums, int k, int left, int right) {  
        int pivot = nums[(left + right) / 2];  
        int i = left, j = right;  
        int t = left;  
        while (t <= j) {  
            if (nums[t] < pivot) {  
                swap(nums, i, t);  
                i++;  
                t++;  
            } else if (nums[t] > pivot) {  
                swap(nums, j, t);  
                j--;  
            } else {  
                t++;  
            }  
        }  
        if (right - j >= k) {  
            return quickSelect(nums, k, j + 1, right);  
        } else if (right - i + 1 >= k) {  
            return pivot;  
        } else {  
            return quickSelect(nums, k - (right - i + 1), left, i - 1);  
        }  
    }  
  
    private void swap(int[] nums, int i, int j) {  
        int temp = nums[i];  
        nums[i] = nums[j];  
        nums[j] = temp;  
    }  
}