数组中第k大的元素

此文为 LeetCode 215. Kth Largest Element in an Array 的题解。

题意

在一个未排序的数组中,寻找第k大的元素。

分析

这个题目其实很简单,但是踩了一些坑,所以记录一下:

首先是用PriorityQueue,然后poll k次即可(5 ms,37.5 MB)。但是属于API caller,略去不表。

其次就是堆排序,很久不写了,第一次把堆排序写成了选择排序,导致时间很长:101 ms,37.2 MB。

后来复习了下堆排序,写出来了(2 ms,36.6 MB):

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
// from https://www.robberphex.com/kth-largest-element-in-an-array/
// 2 ms,36.6 MB
class Solution {
// 调整current子树,将其整理成最小堆
private void adjust(int[] nums, int current, int length) {
// 将堆顶取出来
int num = nums[current];
// 将子树中比顶大的数依次弹上来
for (int i = current * 2 + 1; i < length; i = 2 * i + 1) {
if (i + 1 < length && nums[i + 1] < nums[i]) {
i++;
}
if (nums[i] < num) {
nums[current] = nums[i];
current = i;
} else {
break;
}
}
// 不能弹,则将该数放到对应的位置
nums[current] = num;
}

public int findKthLargest(int[] nums, int k) {
for (int i = (nums.length - 1) / 2; i >= 0; i--) {
// 从第一个非叶子结点开始,构建堆
adjust(nums, i, nums.length);
}
for (int i = 1; i <= nums.length - k+1; i++) {
// 将堆顶交换到数组尾部
int tmp = nums[nums.length - i];
nums[nums.length - i] = nums[0];
nums[0] = tmp;
// 然后继续构建堆
adjust(nums, 0, nums.length - i);
}

return nums[k - 1];
}
}

当然,很容易想到,最优解就是快速选择算法:快速排序划分区间的时候,忽略所有不包含k的区间即可(1 ms,37.3 MB):

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
// from https://www.robberphex.com/kth-largest-element-in-an-array/
// 1 ms,37.3 MB
class Solution {
public int findKthLargest(int[] nums, int k) {
int high = nums.length - 1;
int low = 0;
while (low < high) {
// 选择的区间
int i = low, j = high;
// 从区间两头向中间推进
int pivot = nums[i + (j - i) / 2];
while (i < j) {
// 跳过符合要求的元素
while (i < j && nums[i] > pivot) {
i++;
}
while (i < j && nums[j] < pivot) {
j--;
}
// 如果元素相等,右侧向左移动一位,确保区间一直在缩小
// 同时,确保i是区间的中点(即pivot所在的位置
if (nums[i] == nums[j]) {
j--;
} else {
// 交换元素
int tmp = nums[i];
nums[i] = nums[j];
nums[j] = tmp;
}
}
if (i < k - 1) {
// k在右侧
low = i + 1;
} else if (i > k - 1) {
// k在左侧
high = i - 1;
} else {
// k已经找到
break;
}
}

return nums[k - 1];
}
}
作者

Robert Lu

发布于

2019-07-22

许可协议

评论