数组中的第K个元素
215. 数组中的第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
说明:
你可以假设 k 总是有效的,且 1 ≤ k ≤ 数组的长度。
函数签名如下:
func findKthLargest(nums []int, k int) int
分析
比较直观易解的问题。值得注意的是基于快排思想的快速选择解法
1.朴素实现,时间复杂度O(nlgn),空间复杂度O(1)
func findKthLargest0(nums []int, k int) int {
sort.Sort(sort.Reverse(sort.IntSlice(nums)))
return nums[k-1]
}
2.使用堆
借助一个小顶堆,将nums里的元素一一入堆,但需要保持堆的大小最多为k,如果超出k,需要把堆顶元素出堆 最后堆顶元素就是结果。
时间复杂度O(nlgk),空间复杂度O(k);实际测试,并不比朴素实现快
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) Push(x interface{}) { *h = append(*h, x.(int)) }
func (h *IntHeap) Pop() interface{} {
x := (*h)[len(*h)-1]
*h = (*h)[:len(*h)-1]
return x
}
func findKthLargest1(nums []int, k int) int {
h := &IntHeap{}
for _, v := range nums {
heap.Push(h, v)
if h.Len() > k {
heap.Pop(h)
}
}
return (*h)[0]
}
3.快速选择
时间复杂度 : 平均情况O(N)
,最坏情况O(N^2)
。空间复杂度 : O(1)
func findKthLargest(nums []int, k int) int {
var quickSelect func(lo, hi int)
quickSelect = func(lo, hi int) {
if lo == hi {
return
}
pivotIndex := lo+rand.Intn(hi-lo) // 在[lo, hi]闭区间选择一个随机元素作为基准元素
nums[pivotIndex], nums[hi] = nums[hi], nums[pivotIndex] // 先把基准元素放到最后
pivot := nums[hi]
j := lo
for i := lo; i < hi; i++ {
if nums[i] >= pivot { // 不小于基准元的数字放到前半部分,小于基准元的数字放到后半部分
nums[i], nums[j] = nums[j], nums[i]
j++
}
}
nums[j], nums[hi] = nums[hi], nums[j] // 基准元素放到左右部分之间
if j+1 == k {
return
}
if j+1 > k {
quickSelect(lo, j-1)
} else {
quickSelect(j+1, hi)
}
}
quickSelect(0, len(nums)-1)
return nums[k-1]
}
也可以一开始随机打乱数组,后边每次选择 right 位置的元素为基准元素:
func findKthLargest(nums []int, k int) int {
rand.Shuffle(len(nums), func(i,j int) {
nums[i], nums[j] = nums[j], nums[i]
})
var help func(lo, hi int)
help = func(lo, hi int) {
if lo == hi {
return
}
pivot := nums[hi]
j := lo
for i := lo; i < hi; i++ {
if nums[i] >= pivot {
nums[i], nums[j] = nums[j], nums[i]
j++
}
}
nums[j], nums[hi] = nums[hi], nums[j]
if j+1 == k {
return
}
if j+1 > k {
help(lo, j-1)
} else {
help(j+1, hi)
}
}
help(0, len(nums)-1)
return nums[k-1]
}
实际测试发现比上边每次都随机选基准元素的实现稍微耗时。
扩展: 973. 最接近原点的 K 个点
难度中等
我们有一个由平面上的点组成的列表 points
。需要从中找出 K
个距离原点 (0, 0)
最近的点。
(这里,平面上两点之间的距离是欧几里德距离。)
你可以按任何顺序返回答案。除了点坐标的顺序之外,答案确保是唯一的。
示例 1:
输入:points = [[1,3],[-2,2]], K = 1
输出:[[-2,2]]
解释:
(1, 3) 和原点之间的距离为 sqrt(10),
(-2, 2) 和原点之间的距离为 sqrt(8),
由于 sqrt(8) < sqrt(10),(-2, 2) 离原点更近。
我们只需要距离原点最近的 K = 1 个点,所以答案就是 [[-2,2]]。
示例 2:
输入:points = [[3,3],[5,-1],[-2,4]], K = 2
输出:[[3,3],[-2,4]]
(答案 [[-2,4],[3,3]] 也会被接受。)
提示:
1 <= K <= points.length <= 10000
-10000 < points[i][0] < 10000
-10000 < points[i][1] < 10000
参考解答
func kClosest(points [][]int, k int) [][]int {
rand.Shuffle(len(points), func(i, j int) {
points[i], points[j] = points[j], points[i]
})
var quickSelect func(lo, hi int)
quickSelect = func(lo, hi int) {
if lo == hi {
return
}
pivot := points[hi]
d := dist(pivot)
j := lo
for i := lo; i < hi; i++ {
if dist(points[i]) <= d {
points[j], points[i] = points[i], points[j]
j++
}
}
points[j], points[hi] = points[hi], points[j]
if j+1 == k {
return
}
if j+1 < k {
quickSelect(j+1, hi)
} else {
quickSelect(lo, j-1)
}
}
quickSelect(0, len(points)-1)
return points[:k]
}
func dist(p []int) int {
return p[0]*p[0]+p[1]*p[1]
}