题目描述
给定整数数组 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 <= 105
-104 <= nums[i] <= 104
解法
方法一:快速选择
快速选择算法是一种在未排序的数组中查找第 k
个最大元素或最小元素的算法。它的基本思想是每次选择一个基准元素,将数组分为两部分,一部分的元素都比基准元素小,另一部分的元素都比基准元素大,然后根据基准元素的位置,决定继续在左边还是右边查找,直到找到第 k
个最大元素。
时间复杂度 $O(n)$,空间复杂度 $O(\log n)$。其中 $n$ 为数组 $\textit{nums}$ 的长度。
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25 | class Solution:
def findKthLargest(self, nums: List[int], k: int) -> int:
def quick_sort(l: int, r: int) -> int:
if l == r:
return nums[l]
i, j = l - 1, r + 1
x = nums[(l + r) >> 1]
while i < j:
while 1:
i += 1
if nums[i] >= x:
break
while 1:
j -= 1
if nums[j] <= x:
break
if i < j:
nums[i], nums[j] = nums[j], nums[i]
if j < k:
return quick_sort(j + 1, r)
return quick_sort(l, j)
n = len(nums)
k = n - k
return quick_sort(0, n - 1)
|
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33 | class Solution {
private int[] nums;
private int k;
public int findKthLargest(int[] nums, int k) {
this.nums = nums;
this.k = nums.length - k;
return quickSort(0, nums.length - 1);
}
private int quickSort(int l, int r) {
if (l == r) {
return nums[l];
}
int i = l - 1, j = r + 1;
int x = nums[(l + r) >>> 1];
while (i < j) {
while (nums[++i] < x) {
}
while (nums[--j] > x) {
}
if (i < j) {
int t = nums[i];
nums[i] = nums[j];
nums[j] = t;
}
}
if (j < k) {
return quickSort(j + 1, r);
}
return quickSort(l, j);
}
}
|
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28 | class Solution {
public:
int findKthLargest(vector<int>& nums, int k) {
int n = nums.size();
k = n - k;
auto quickSort = [&](auto&& quickSort, int l, int r) -> int {
if (l == r) {
return nums[l];
}
int i = l - 1, j = r + 1;
int x = nums[(l + r) >> 1];
while (i < j) {
while (nums[++i] < x) {
}
while (nums[--j] > x) {
}
if (i < j) {
swap(nums[i], nums[j]);
}
}
if (j < k) {
return quickSort(quickSort, j + 1, r);
}
return quickSort(quickSort, l, j);
};
return quickSort(quickSort, 0, n - 1);
}
};
|
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33 | func findKthLargest(nums []int, k int) int {
k = len(nums) - k
var quickSort func(l, r int) int
quickSort = func(l, r int) int {
if l == r {
return nums[l]
}
i, j := l-1, r+1
x := nums[(l+r)>>1]
for i < j {
for {
i++
if nums[i] >= x {
break
}
}
for {
j--
if nums[j] <= x {
break
}
}
if i < j {
nums[i], nums[j] = nums[j], nums[i]
}
}
if j < k {
return quickSort(j+1, r)
}
return quickSort(l, j)
}
return quickSort(0, len(nums)-1)
}
|
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23 | function findKthLargest(nums: number[], k: number): number {
const n = nums.length;
k = n - k;
const quickSort = (l: number, r: number): number => {
if (l === r) {
return nums[l];
}
let [i, j] = [l - 1, r + 1];
const x = nums[(l + r) >> 1];
while (i < j) {
while (nums[++i] < x);
while (nums[--j] > x);
if (i < j) {
[nums[i], nums[j]] = [nums[j], nums[i]];
}
}
if (j < k) {
return quickSort(j + 1, r);
}
return quickSort(l, j);
};
return quickSort(0, n - 1);
}
|
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39 | impl Solution {
pub fn find_kth_largest(mut nums: Vec<i32>, k: i32) -> i32 {
let len = nums.len();
let k = len - k as usize;
Self::quick_sort(&mut nums, 0, len - 1, k)
}
fn quick_sort(nums: &mut Vec<i32>, l: usize, r: usize, k: usize) -> i32 {
if l == r {
return nums[l];
}
let (mut i, mut j) = (l as isize - 1, r as isize + 1);
let x = nums[(l + r) / 2];
while i < j {
i += 1;
while nums[i as usize] < x {
i += 1;
}
j -= 1;
while nums[j as usize] > x {
j -= 1;
}
if i < j {
nums.swap(i as usize, j as usize);
}
}
let j = j as usize;
if j < k {
Self::quick_sort(nums, j + 1, r, k)
} else {
Self::quick_sort(nums, l, j, k)
}
}
}
|
方法二:优先队列(小根堆)
我们可以维护一个大小为 $k$ 的小根堆 $\textit{minQ}$,然后遍历数组 $\textit{nums}$,将数组中的元素依次加入到小根堆中,当小根堆的大小超过 $k$ 时,我们将堆顶元素弹出,这样最终小根堆中的 $k$ 个元素就是数组中的 $k$ 个最大元素,堆顶元素就是第 $k$ 个最大元素。
时间复杂度 $O(n\log k)$,空间复杂度 $O(k)$。其中 $n$ 为数组 $\textit{nums}$ 的长度。
| class Solution:
def findKthLargest(self, nums: List[int], k: int) -> int:
return nlargest(k, nums)[-1]
|
1
2
3
4
5
6
7
8
9
10
11
12 | class Solution {
public int findKthLargest(int[] nums, int k) {
PriorityQueue<Integer> minQ = new PriorityQueue<>();
for (int x : nums) {
minQ.offer(x);
if (minQ.size() > k) {
minQ.poll();
}
}
return minQ.peek();
}
}
|
1
2
3
4
5
6
7
8
9
10
11
12
13 | class Solution {
public:
int findKthLargest(vector<int>& nums, int k) {
priority_queue<int, vector<int>, greater<int>> minQ;
for (int x : nums) {
minQ.push(x);
if (minQ.size() > k) {
minQ.pop();
}
}
return minQ.top();
}
};
|
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21 | func findKthLargest(nums []int, k int) int {
minQ := hp{}
for _, x := range nums {
heap.Push(&minQ, x)
if minQ.Len() > k {
heap.Pop(&minQ)
}
}
return minQ.IntSlice[0]
}
type hp struct{ sort.IntSlice }
func (h hp) Less(i, j int) bool { return h.IntSlice[i] < h.IntSlice[j] }
func (h *hp) Push(v any) { h.IntSlice = append(h.IntSlice, v.(int)) }
func (h *hp) Pop() any {
a := h.IntSlice
v := a[len(a)-1]
h.IntSlice = a[:len(a)-1]
return v
}
|
| function findKthLargest(nums: number[], k: number): number {
const minQ = new MinPriorityQueue();
for (const x of nums) {
minQ.enqueue(x);
if (minQ.size() > k) {
minQ.dequeue();
}
}
return minQ.front().element;
}
|
1
2
3
4
5
6
7
8
9
10
11
12
13
14 | use std::collections::BinaryHeap;
impl Solution {
pub fn find_kth_largest(nums: Vec<i32>, k: i32) -> i32 {
let mut minQ = BinaryHeap::new();
for &x in nums.iter() {
minQ.push(-x);
if minQ.len() > k as usize {
minQ.pop();
}
}
-minQ.peek().unwrap()
}
}
|
方法三:计数排序
我们可以使用计数排序的思想,统计数组 $\textit{nums}$ 中每个元素出现的次数,记录在哈希表 $\textit{cnt}$ 中,然后从大到小遍历元素 $i$,每次减去出现的次数 $\textit{cnt}[i]$,直到 $k$ 小于等于 $0$,此时的元素 $i$ 就是数组中的第 $k$ 个最大元素。
时间复杂度 $O(n + m)$,空间复杂度 $O(n)$。其中 $n$ 为数组 $\textit{nums}$ 的长度,而 $m$ 为数组 $\textit{nums}$ 中元素的最大值。
| class Solution:
def findKthLargest(self, nums: List[int], k: int) -> int:
cnt = Counter(nums)
for i in count(max(cnt), -1):
k -= cnt[i]
if k <= 0:
return i
|
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16 | class Solution {
public int findKthLargest(int[] nums, int k) {
Map<Integer, Integer> cnt = new HashMap<>(nums.length);
int m = Integer.MIN_VALUE;
for (int x : nums) {
m = Math.max(m, x);
cnt.merge(x, 1, Integer::sum);
}
for (int i = m;; --i) {
k -= cnt.getOrDefault(i, 0);
if (k <= 0) {
return i;
}
}
}
}
|
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17 | class Solution {
public:
int findKthLargest(vector<int>& nums, int k) {
unordered_map<int, int> cnt;
int m = INT_MIN;
for (int x : nums) {
++cnt[x];
m = max(m, x);
}
for (int i = m;; --i) {
k -= cnt[i];
if (k <= 0) {
return i;
}
}
}
};
|
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15 | func findKthLargest(nums []int, k int) int {
cnt := map[int]int{}
m := -(1 << 30)
for _, x := range nums {
cnt[x]++
m = max(m, x)
}
for i := m; ; i-- {
k -= cnt[i]
if k <= 0 {
return i
}
}
}
|
1
2
3
4
5
6
7
8
9
10
11
12
13 | function findKthLargest(nums: number[], k: number): number {
const cnt: Record<number, number> = {};
for (const x of nums) {
cnt[x] = (cnt[x] || 0) + 1;
}
const m = Math.max(...nums);
for (let i = m; ; --i) {
k -= cnt[i] || 0;
if (k <= 0) {
return i;
}
}
}
|
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27 | use std::collections::HashMap;
impl Solution {
pub fn find_kth_largest(nums: Vec<i32>, k: i32) -> i32 {
let mut cnt = HashMap::new();
let mut m = i32::MIN;
for &x in &nums {
*cnt.entry(x).or_insert(0) += 1;
if x > m {
m = x;
}
}
let mut k = k;
for i in (i32::MIN..=m).rev() {
if let Some(&count) = cnt.get(&i) {
k -= count;
if k <= 0 {
return i;
}
}
}
unreachable!();
}
}
|