本篇是来自 Swift 算法学院的翻译的一篇文章,Swift 算法学院 致力于使用 Swift 实现各种算法,对想学习算法或者复习算法的同学非常有帮助,讲解思路非常清楚,每一篇都有详细的例子解释。 更多翻译的文章还可以查看这里。
第K大元素
给定一个数组 a
,写一个算法找出第K大的元素。
比如在 第一大 的元素是最大元素。如果数组有 n 个元素,第 n 大 元素为最小值,中间最大为 第n/2 大值。
原始方案
下面的算法是半原生的。它的时间复杂度是 O(n log n),因为它需要先排序,因此需要额外的 O(n) 的空间。
func kthLargest(a: [Int], k: Int) -> Int? {
let len = a.count
if k > 0 && k <= len {
let sorted = a.sorted()
return sorted[len - k]
} else {
return nil
}
}
kthLargest()
函数有两个参数,整数型数组 a
和 k
用来表示第 k 大的元素。
举例说明一下这个算法的原理,假定 k = 4
, 数组如下:
[ 7, 92, 23, 9, -1, 0, 11, 6 ]
最开始无法直接找到第 k 大的元素,但是排序后就非常简单了,排序后如下:
[ -1, 0, 6, 7, 9, 11, 23, 92 ]
现在只需要取 a.count - k
对应的值:
a[a.count - k] = a[8 - 4] = a[4] = 9
当然如果需要找 第 k 小 的值时用 a[k-1]
即可
更快的算法
这个算法借鉴了 二分查找 和 快排 的思想,时间复杂度为 O(n)
不断调用二分查找将数组分割成一半又一半,快速的缩小查询的值的范围。
快速排序也分割数组,把小于轴值的移至左边,所有大于轴值的移至右边。经过某个轴值分区后,轴值所在的位置就是排序后最终位置。可以利用这一点来提高算法。
下面介绍如何工作:随机选一个值作为轴值进行分区,像二分查找一样继续对左右分区进行处理,直到恰好一个轴值是在 k-th
位置。
举个例子说明一下,在下面的数组中找 第 4
大的元素:
[ 7, 92, 23, 9, -1, 0, 11, 6 ]
该算法对查找第k小值也是很简单的,来让我们试试查找k = 4
的最小值。
我们不用先对数组排序,随机选一个值比如 11
作为轴值进行分区,结果如下:
[ 7, 9, -1, 0, 6, 11, 92, 23 ]
<------ smaller larger -->
根据结果,比 11
小的值在左边,大的值在右边。11
在它的最终位置上,索引值为 5 , 因此第 4 小的值肯定是在左边的位置可以忽略其他的部分:
[ 7, 9, -1, 0, 6, x, x, x ]
再随机选一个轴值比如 6
将数组分区,结果如下:
[ -1, 0, 6, 9, 7, x, x, x ]
轴值 6
的索引值为 2,显然第 4
大的值在右边分区,可以忽略左边的分区了:
[ x, x, x, 9, 7, x, x, x ]
重复以上操作后如下:
[ x, x, x, 7, 9, x, x, x ]
轴值 9
的索引值为 4,而且这正是要查找的!可以看到我们不需要对数组排序,用很少的步数就能实现。
实现方法如下:
public func randomizedSelect<T: Comparable>(_ array: [T], order k: Int) -> T {
var a = array
func randomPivot<T: Comparable>(_ a: inout [T], _ low: Int, _ high: Int) -> T {
let pivotIndex = random(min: low, max: high)
a.swapAt(pivotIndex, high)
return a[high]
}
func randomizedPartition<T: Comparable>(_ a: inout [T], _ low: Int, _ high: Int) -> Int {
let pivot = randomPivot(&a, low, high)
var i = low
for j in low..<high {
if a[j] <= pivot {
a.swapAt(i, j)
i += 1
}
}
a.swapAt(i, high)
return i
}
func randomizedSelect<T: Comparable>(_ a: inout [T], _ low: Int, _ high: Int, _ k: Int) -> T {
if low < high {
let p = randomizedPartition(&a, low, high)
if k == p {
return a[p]
} else if k < p {
return randomizedSelect(&a, low, p - 1, k)
} else {
return randomizedSelect(&a, p + 1, high, k)
}
} else {
return a[low]
}
}
precondition(a.count > 0)
return randomizedSelect(&a, 0, a.count - 1, k)
}
为了提高可读性,这个函数分成三个内部函数:
randomPivot()
随机选取一个数字,然后放在当前分区的最后一个位置(这是Lomuto 分区方式所规定的,更多介绍请看快排)randomizedPartition()
是快排中 Lomuto 分区方法。当完成后,随机轴值在的位置就是排序后的最终位置。返回轴值所在的位置。randomizedSelect()
做所有的脏活累活。先调用分区函数,后决定再做什么。如果轴值索引值等于 k ,那么该值正是查找值,完成查找。如果k
比该索引值小,那么查找值一定在左边分区,递归调用就可以了,否则就肯定是在右边分区中。
非常😎,是不是? 快排的期望复杂度为 o(n log n), 但是因为只把数组分成越来越小的分区,randomizedSelect()
的时间复杂度为 O(n)。
注意:该函数式计算数组中 第k 小元素,
k
是从 0 开始的。如果需要第k
大元素,应调用a.count - k
。
作者 Daniel Speiser 修改 Matthijs Hollemans 译者KeithMorning