给定整数数组 nums 和整数 k,请返回数组中第 k 个最大的元素。
请注意,你需要找的是数组排序后的第 k 个最大的元素,而不是第 k 个不同的元素。
示例 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 <= 104
-104 <= nums[i] <= 104
最大堆
手写最大堆,涉及到两个方面,一个是初始化建立最大堆,一个是删除堆顶元素并调整最大堆。
初始化最大堆
思路: 考虑定义一个函数
maxHeapify(int[] nums, int idx, int heapSize)
我们预想 以nums[idx]为根节点的左子树已经是一个最大堆,以nums[idx]为根节点的右子树也已经是一个最大堆。执行maxHeapify函数后,我们想以nums[idx]为根节点的树也是一个最大堆,maxHeapify函数就是做这个事情的。
如何实现这个函数呢?
1 . idx的左子节点为 left = nums[2*idx + 1]
, 右子节点right = nums[2*idx + 2]
。
如果left 和right都不大于 target = nums[idx]
,则不用做任何处理,idx为根节点的树就是一个最大堆
2 . 如果left,right中有一个或者两个大于target, 找出最大 的那个,假设是left(right分析类似), 将其与target交换,则 target, left, right 仅仅这三个元素连起来看是一个最大堆。
right子树不用调整(因为right依然小于target),但是left子树却可能需要再迭代处理,因为left节点的值与target发生了交换,变小了,可能left子树不再是最大堆了。但是left节点的左子树和右子树却依然是最大堆(这个是在预想之中的), 所以我们需要对left节点依然做 maxHeapity
调整。
我们发现,这与最开始考虑的调整target
子树是一个问题,规模变小,这就是递归。我们在递归时还需要使我们的预想成立(以nums[idx]为根节点的左子树已经是一个最大堆,以nums[idx]为根节点的右子树也已经是一个最大堆),这就要求我们在初始化最大堆时,从叶子的上一层开始构建最大堆,构建好之后,再逐一上一层构建最大堆,最后到根节点构建最大堆。
// idx的左子树,右子树已经是最大堆,将idx节点加进去调整,再构建最大堆
public void maxHeapify(int[] nums, int idx, int heapSize){
int largestIdx = idx;
int largest = nums[idx];
int leftIdx = 2*idx+1;
int left = nums[leftIdx];
int rightIdx = 2*idx+2;
int right = nums[rightIdx];
// 左节点比父节点大
if(leftIdx < heapSize && left > largest){
largest = left;
largestIdx = leftIdx;
}
// 右节点比父节点更大
if(rightIdx < heapSize && right > largest){
largest = right;
largestIdx = rightIdx
}
// 存在左右子节点,比父亲大的情况
if(largestIdx != idx){
// 交换
swap(nums, idx, largestIdx);
// 递归调整
maxHeapify(nums, largestIdx, heapSize);
}
}
public void swap(int[] a, int i, int j) {
int temp = a[i];
a[i] = a[j];
a[j] = temp;
}
ok,刚才说到初始化最大堆需要,从叶子节点上一层逐层调整堆,所以还会有以下for循环代码
for(int idx = heapSize/2; idx>=0; idx--){
maxHeapify(nums, idx, heapSize);
}
问: 为什么从idx = heapSize/2开始构建最大堆,再向上层递进调整最大堆。
注意到idx = heapSize/2时,左孩子在数组下标是 2*idx+1, 其可能等于heapSize+1(当heapSize为偶数时), 也可能等于heapSize(当heapSize为奇数时), 都会正好超出数组边界。右孩子节点必定超出数组边界。
int largest = nums[idx];
int leftIdx = 2*idx+1;
int left = nums[leftIdx];
上面这种写法需要bugfix,因为求left时可能会数组越界。
另外, idx可以从heapSize/2-1开始算
for(int idx = heapSize/2 - 1; idx>=0; idx–){
maxHeapify(nums, idx, heapSize);
}
删除堆顶部元素,调整堆
当k=1,即求第一大元素时,不用删除堆顶元素,堆顶元素即为第一大元素。
当k=2,即求第二大元素,删除一次堆顶元素,将最后一个节点放置到堆顶,调整最大堆,堆顶元素即为最大值。
…
所以,寻找第K大元素,就是删除并调整最大堆K-1次,获取堆顶元素。
for(int i = 1; i<=K-1;i++){
nums[0] = nums[heapSize-1];
heapSize--;
maxHeapify(nums, 0, heapSize);
}