215. 第 K 个最大元素
题目描述
输入整数数组 nums 和整数 k。把数组从大到小排好,返回排在第 k 位的元素。
重复的值要分别计数,所以要找的不是第 k 个互不相同的值。
你需要设计时间复杂度为 O(n) 的算法。
示例 1
输入:nums = [3,2,1,5,6,4], k = 2 输出:5
解释:从大到小为 6、5、4…,第 2 大是 5。
示例 2
输入:nums = [3,2,3,1,2,4,5,5,6], k = 4 输出:4
解释:从大到小为 6、5、5、4…,第 4 大是 4。
数据范围
1 <= k <= nums.length <= 10^5-10^4 <= nums[i] <= 10^4
解题思路:大小为 k 的小根堆
维护容量为 k 的最小堆。
为什么正确
堆中保存当前最大的 k 个元素。
复杂度分析
时间 O(n log k),额外空间 O(k)。
易错点
注意边界条件并在回溯或弹出后恢复状态。
参考代码
Python、Java、C++、Go、JavaScript、TypeScript 六种语言,逻辑一致。
import heapq
class Solution:
def findKthLargest(self,a,k):
h=[]
for x in a:heapq.heappush(h,x);heapq.heappop(h) if len(h)>k else None
return h[0]class Solution {
public int findKthLargest(int[] a, int k) {
PriorityQueue<Integer> h = new PriorityQueue<>();
for (int x : a) { h.offer(x); if (h.size() > k) h.poll(); }
return h.peek();
}
}class Solution {
public:
int findKthLargest(vector<int>& a, int k) {
priority_queue<int, vector<int>, greater<int>> h;
for (int x : a) { h.push(x); if (h.size() > k) h.pop(); }
return h.top();
}
};type minHeap []int
func (h minHeap) Len() int { return len(h) }
func (h minHeap) Less(i, j int) bool { return h[i] < h[j] }
func (h minHeap) Swap(i, j int) { h[i], h[j] = h[j], h[i] }
func (h *minHeap) Push(x any) { *h = append(*h, x.(int)) }
func (h *minHeap) Pop() any { old := *h; x := old[len(old)-1]; *h = old[:len(old)-1]; return x }
func findKthLargest(a []int, k int) int {
h := &minHeap{}
for _, x := range a {
heap.Push(h, x)
if h.Len() > k { heap.Pop(h) }
}
return (*h)[0]
}const heappush = (h, x) => {
h.push(x);
for (let i = h.length - 1, p; i > 0 && h[p = (i - 1) >> 1] > h[i]; i = p) [h[p], h[i]] = [h[i], h[p]];
};
const heappop = h => {
const top = h[0], last = h.pop();
if (h.length) {
h[0] = last;
for (let i = 0; ;) {
const l = 2 * i + 1, r = l + 1; let s = i;
if (l < h.length && h[l] < h[s]) s = l;
if (r < h.length && h[r] < h[s]) s = r;
if (s === i) break;
[h[s], h[i]] = [h[i], h[s]]; i = s;
}
}
return top;
};
var findKthLargest = function(a, k) {
const h = [];
for (const x of a) { heappush(h, x); if (h.length > k) heappop(h); }
return h[0];
};const heappush = (h: number[], x: number): void => {
h.push(x);
for (let i = h.length - 1, p: number; i > 0 && h[p = (i - 1) >> 1] > h[i]; i = p) [h[p], h[i]] = [h[i], h[p]];
};
const heappop = (h: number[]): number => {
const top = h[0], last = h.pop()!;
if (h.length) {
h[0] = last;
for (let i = 0; ;) {
const l = 2 * i + 1, r = l + 1; let s = i;
if (l < h.length && h[l] < h[s]) s = l;
if (r < h.length && h[r] < h[s]) s = r;
if (s === i) break;
[h[s], h[i]] = [h[i], h[s]]; i = s;
}
}
return top;
};
function findKthLargest(a: number[], k: number): number {
const h: number[] = [];
for (const x of a) { heappush(h, x); if (h.length > k) heappop(h); }
return h[0];
};其他解法
堆解法稳定 O(n log k),不修改原数组,也适合数据流;快速选择每次只递归一侧,平均 O(n),但最坏会退化,并且需要一次性拿到整个数组。
解法 2:快速选择 时间 平均 O(n),最坏 O(n²),空间 O(1)
第 k 大就是升序排列后下标为 n-k 的元素。借用快速排序的划分:取区间中点的值作为基准,用 Hoare 划分把区间分成左边都 ≤ 基准、右边都 ≥ 基准的两段,分界为 j。目标下标在哪一段,就只在那一段继续划分,另一段直接丢弃,直到区间只剩一个元素。
平均线性时间、原地进行,比堆快。最坏情况退化到 O(n²);会修改原数组;划分的边界条件容易写错。Hoare 划分在大量重复元素时仍能均匀划分。
from typing import List
class Solution:
def findKthLargest(self, nums: List[int], k: int) -> int:
target = len(nums) - k
l, r = 0, len(nums) - 1
while l < r:
pivot = nums[(l + r) // 2]
i, j = l - 1, r + 1
while True:
i += 1
while nums[i] < pivot:
i += 1
j -= 1
while nums[j] > pivot:
j -= 1
if i >= j:
break
nums[i], nums[j] = nums[j], nums[i]
if target <= j:
r = j
else:
l = j + 1
return nums[l]class Solution {
public int findKthLargest(int[] nums, int k) {
int target = nums.length - k;
int l = 0, r = nums.length - 1;
while (l < r) {
int pivot = nums[l + (r - l) / 2];
int i = l - 1, j = r + 1;
while (true) {
do i++; while (nums[i] < pivot);
do j--; while (nums[j] > pivot);
if (i >= j) break;
int t = nums[i]; nums[i] = nums[j]; nums[j] = t;
}
if (target <= j) r = j;
else l = j + 1;
}
return nums[l];
}
}class Solution {
public:
int findKthLargest(vector<int>& nums, int k) {
int target = nums.size() - k;
int l = 0, r = nums.size() - 1;
while (l < r) {
int pivot = nums[l + (r - l) / 2];
int i = l - 1, j = r + 1;
while (true) {
do i++; while (nums[i] < pivot);
do j--; while (nums[j] > pivot);
if (i >= j) break;
swap(nums[i], nums[j]);
}
if (target <= j) r = j;
else l = j + 1;
}
return nums[l];
}
};func findKthLargest(nums []int, k int) int {
target := len(nums) - k
l, r := 0, len(nums)-1
for l < r {
pivot := nums[l+(r-l)/2]
i, j := l-1, r+1
for {
i++
for nums[i] < pivot {
i++
}
j--
for nums[j] > pivot {
j--
}
if i >= j {
break
}
nums[i], nums[j] = nums[j], nums[i]
}
if target <= j {
r = j
} else {
l = j + 1
}
}
return nums[l]
}var findKthLargest = function(nums, k) {
const target = nums.length - k;
let l = 0, r = nums.length - 1;
while (l < r) {
const pivot = nums[(l + r) >> 1];
let i = l - 1, j = r + 1;
while (true) {
do i++; while (nums[i] < pivot);
do j--; while (nums[j] > pivot);
if (i >= j) break;
[nums[i], nums[j]] = [nums[j], nums[i]];
}
if (target <= j) r = j;
else l = j + 1;
}
return nums[l];
};function findKthLargest(nums: number[], k: number): number {
const target = nums.length - k;
let l = 0, r = nums.length - 1;
while (l < r) {
const pivot = nums[(l + r) >> 1];
let i = l - 1, j = r + 1;
while (true) {
do i++; while (nums[i] < pivot);
do j--; while (nums[j] > pivot);
if (i >= j) break;
[nums[i], nums[j]] = [nums[j], nums[i]];
}
if (target <= j) r = j;
else l = j + 1;
}
return nums[l];
};