题目:
给定整数数组 nums
和整数 k
,请返回数组中第 k
个最大的元素。
请注意,你需要找的是数组排序后的第 k
个最大的元素,而不是第 k
个不同的元素。
1 <= k <= nums.length <= 104
-104 <= nums[i] <= 104
本算法题目有两种比较好的方法,分别为使用小顶堆
和使用快速选择
。
题解一:
首先我们使用小顶堆
的方法求解,先贴代码(Go语言):
type IntHeap []int
func (h IntHeap) Len() int { return len(h) }
func (h IntHeap) Less(i, j int) bool { return h[i] < h[j] }
func (h IntHeap) Swap(i, j int) { h[i], h[j] = h[j], h[i] }
func (h IntHeap) Top() int {
return h[0]
}
func (h *IntHeap) Push(x interface{}) {
// Push and Pop use pointer receivers because they modify the slice's length,
// not just its contents.
*h = append(*h, x.(int))
}
func (h *IntHeap) Pop() interface{} {
old := *h
n := len(old)
x := old[n-1]
*h = old[0 : n-1]
return x
}
func findKthLargest(nums []int, k int) int {
h := &IntHeap{}
heap.Init(h)
for _, num := range nums {
if h.Len() < k {
heap.Push(h, num)
} else if (*h)[0] < num {
heap.Pop(h)
heap.Push(h, num)
}
}
return heap.Pop(h).(int)
}
首先我们确定循环不变量:h
中始终存储了最多k
个数组元素,且是目前遍历过的最大的k
个元素,且其中最小的元素为(*h)[0]
。
h
是一个小顶堆,则h
能够保证,加入其中的元素的最小的元素为(*h)[0]
。(小顶堆的概念可自行Google)
初始化: h
中没有元素,不变式成立。
保持: 在迭代过程中,根据h
的元素数量,分为如下两种情况:
- 如果
h
中元素的个数小于k
,直接将本次迭代的数组元素num
加入到h
中。 - 如果
h
中元素的个数大于等于k
,则需要比较当前h
中,最小的元素是否大于num
,若是,忽略num
的处理,否则,将其替换为num
。
以上两种情况的结果均可以保证不变式成立。
终止: nums
中的所有元素均被遍历后算法终止,不变式成立。
最终h
中的最小元素(*h)[0]
即为数组nums
的第K
个最大元素。
复杂度分析:
时间复杂度:O(nlogn)
,建堆的时间代价是 O(n)
,删除的总代价是O(klogn)
,因为 k < nk < n
,故渐进时间复杂为O(n + klogn) = O(nlogn)
。
空间复杂度:O(logn)
,即递归使用栈空间的空间代价。
题解二:
我们同样可以使用快速选择
的方法求解,先贴代码(Go语言):
func findKthLargest(nums []int, k int) int {
rand.Seed(time.Now().UnixNano())
return quickSelect(nums, 0, len(nums)-1, k)
}
func quickSelect(nums []int, l int, r int, k int) int {
m := randomPartition(nums, l, r)
if r-m+1 == k {
return nums[m]
} else if r-m+1 > k {
return quickSelect(nums, m+1, r, k)
} else {
return quickSelect(nums, l, m-1, k-(r-m+1))
}
}
func randomPartition(nums []int, l, r int) int {
i := rand.Int()%(r-l+1) + l
nums[i], nums[r] = nums[r], nums[i]
return partition(nums, l, r)
}
func partition(nums []int, l int, r int) int {
key := nums[r]
//all in [l,i) < key
//all in [i,j] > key
i := l
j := l
for j < r {
if nums[j] < key {
nums[i], nums[j] = nums[j], nums[i]
i++
}
j++
}
nums[i], nums[r] = nums[r], nums[i]
return i
}
快速选择
与快速排序
很类似,使用了相同的分治思想,(快速排序
算法可参考这篇文章)将快速排序
的处理阶段稍加修改即可。同时我们这里的快速选择
算法使用了随机采样
的方法,提高了算法的性能。
复杂度分析:
时间复杂度:O(n)
,如上文所述,证明过程可以参考「《算法导论》9.2:期望为线性的选择算法」。
空间复杂度:O(log n)
,递归使用栈空间的空间代价的期望为 O(logn)
。
两种解法的对比:
快速选择
算法时间复杂度更低,且不需要额外的存储空间。但是必须预先知道数组的总长度,所以适合对定长数组的计算。堆
需要额外的k
个存储空间,且算法时间复杂度相对高一些。但是不需要预先知道待处理数组的规模,所以适合对流数据进行处理。