LeetCode215 数组中的第K个最大元素(含堆解法实现)

链接: https://leetcode-cn.com/problems/kth-largest-element-in-an-array/

题面

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

解法

经典老番
《茴有几种写法》级别的题目

这里主要复习了一下最大堆最小堆的实现,以这道题作为背景
着重体会关于下沉时的操作 细节见代码

后面还有快排partition以及优先队列解法

代码

堆解法

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
40
41
42
43
44
45
46
47
48
49
50
class Solution {
public:
vector<int> heap;
// 实现一个最小堆
int top() {
return heap[1];
}

void swim(int k) {
while(k > 1 && heap[k / 2] > heap[k]) {
swap(heap[k / 2], heap[k]);
k /= 2;
}
}
void push(int num) {
heap.push_back(num);
swim(heap.size() - 1);
}
void sink(int k) {
while(k * 2 < heap.size()) {
// 如果当前结点有子结点
int pos = k * 2;
if (pos + 1 < heap.size() && heap[pos + 1] < heap[pos]) {
pos++; // 移到右子结点
}
if (heap[k] <= heap[pos]) { // 如果两个子结点都比父结点要大
break; // sink结束
}
swap(heap[k], heap[pos]); // 交换子结点和父结点
k = pos;
}
}
void pop() {
heap[1] = heap.back();
heap.pop_back();
sink(1);
}

int findKthLargest(vector<int>& nums, int k) {
heap.push_back(-1);
for (int i: nums) {
push(i);
if (heap.size() > k + 1) {
pop();
}
}
return top();

}
};

集大成者

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
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
class Solution {
public:
void arrangeK(vector<int>& nums, int l, int r, int k) {
if (l > r) {
return;
}
int key = nums[l], a = l, b = r;
while (a < b) {
while (a < b && nums[b] >= key) {
b--;
}
while(a < b && nums[a] <= key) {
a++;
}
swap(nums[a], nums[b]);
}
swap(nums[l], nums[a]);
// 最大的k个元素已置于右端
if (r - a + 1 == k) {
return;
}
else if (r - a + 1 < k) {
arrangeK(nums, l, a - 1, k - (r - a + 1));
}
else {
arrangeK(nums, a, r, k);
}
}
int findKdc(vector<int>& nums, int k) {
arrangeK(nums, 0, nums.size() - 1, k);
int res = nums.back(), len = nums.size();
for(int i = 0; i < k; i++) {
res = min(res, nums[len - 1 - i]);
}
return res;
}
int quickSel(vector<int>& nums, int l, int r) {
int key = nums[l];
int a = l, b = r;
while(a < b) {
while(a < b && nums[b] >= key) {
b--;
}
while(a < b && nums[a] <= key) {
a++;
}
swap(nums[a], nums[b]);
}
swap(nums[a], nums[l]);
return a;

}

int solve(vector<int>& nums, int k) {
int l = 0, r = nums.size() - 1, target = nums.size() - k;
// 找第k大的数也就是要有len - k个比其小的数
while (l < r) {
int mid = quickSel(nums, l, r);
if (mid == target) {
return nums[mid];
}
else if (mid < target) {
l = mid + 1;
}
else {
r = mid - 1;
}
}
return nums[l];
}
int findWithPQ(vector<int>& nums, int k) {
priority_queue<int, vector<int>, greater<int> > pq;
for (int i: nums) {
pq.push(i);
if (pq.size() > k) {
pq.pop();
}
}
return pq.top();
}
int findKthLargest(vector<int>& nums, int k) {
return findWithPQ(nums, k);
}
};

partition精简版

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
class Solution {
public:
int partition(vector<int>& a, int l, int r) {
swap(a[l], a[r]);
int small = l;
for (int i = l; i < r; i++) {
if (a[i] < a[r]) {
swap(a[i], a[small]);
small++;
}
}
swap(a[small], a[r]);
return small;
}
void findk(vector<int>& nums, int l, int r, int k) {
int index = partition(nums, l, r);
while (index != k) {
if (index < k) {
l = index + 1;
}
else {
r = index - 1;
}
int pivot = rand() % (r - l + 1) + l;
swap(nums[l], nums[pivot]);
index = partition(nums, l, r);
}
}
int findKthLargest(vector<int>& nums, int k) {
int n = nums.size();
findk(nums, 0, n - 1, n - k);
return nums[n - k];
}
};