浅谈topk问题——文中给出两种C++实现

TopK问题是要求设计一种算法————能在 大量无序数据中 快速找到 第k大 或者 第k小 的元素。

本文暂时不考虑超大文件。那么求取第k个元素,一般有四种思路, - 全部排序,选取第k个 - 冒泡排序,(冒泡一次就确定一个数字,不需要排序全部) - 改进的堆排序 - 改进的快速排序

对于全部排序和冒泡排序,上一篇文章已经讲过了《一文实现基本排序算法》,本文不再赘述。

改进的堆排序

时间复杂度: O(n * log(k))

如果需要求“第k大”的元素,那么需要建立大小为k的“小顶堆”。此时,因为堆里总共k个元素,堆顶是最小元素,所以它就是“第k大”。以“第k大”为例子,具体思路,

假设所有数据长度为length(>= k); - 使用前k个元素,创建一个大小为k的小顶堆。 - 从第k+1个元素开始,将后续元素挨个与堆顶元素比较,如果“大于“堆顶元素,则将该元素赋值为堆顶。调整堆,使其满足小顶堆。

算法实现,[https://github.com/qingdujun/algorithm/blob/master/heap-sort-topk.cpp]

template<typename T, typename F>
void adjust_heap(T* begin, int parent, int length, F cmp) {
T tmp = *(begin + parent);
for (int child = 2 * parent + 1; child < length; child = 2 * parent + 1) {
if (child+1 < length && cmp(*(begin+child+1), *(begin+child)))
++child;
if (cmp(tmp, *(begin+child)))
break;
*(begin+parent) = *(begin+child);
parent = child;
}
*(begin+parent) = tmp;
}

template<typename T, typename F>
T heap_sort_topk(T* begin, T* end, int k, F cmp) {
if (begin >= end || begin == nullptr || end == nullptr || cmp == nullptr || k <= 0)
throw std::invalid_argument("invalid arg");
int length = end - begin;
for (int parent = k / 2 - 1; parent >= 0; --parent) {
adjust_heap(begin, parent, length, cmp);
}
for (int next = k; next < length; ++next) {
if ( cmp(*begin, *(begin+next)) ) {
std::swap(*begin, *(begin+next));
adjust_heap(begin, 0, k, cmp);
}
}
return *begin;
}

改进的快速排序

时间复杂度: O(n)

你已经知道,快速排序每次划分会最终确定一个元素的位置。利用这个特性,可以快速的求出“第k小”元素。具体思路如下,

假设所有数据长度为length(>= k); - 对数据进行一次划分,最终正确排序位置为pos。 - 如果,pos == k,那么直接返回,找到正确结果 - 如果,pos > k,那么应该取左边,继续划分 - 如果,pos < k, 那么说明应该取右边继续划分,但是,此时的k应该更改为 k -= pos

当然,还有一种其他的实现方法,可以参考《BFPRT算法:时间复杂度O(n)求第k小的数字(分治算法+快排)》

分治算法边界确定要十分注意,我原来想按照STL中的[begin, end)确定边界,结果发现有时候运行会出错(比如*right),所以以下边界统一为[begin, end]。

算法实现,[https://github.com/qingdujun/algorithm/blob/master/quick_sort_topk.cpp]

template<typename T, typename F>
T quick_sort_topk(T* begin, T* end, int k, F cmp) {//[begin, end]
if (begin == nullptr || end == nullptr || cmp == nullptr || k <= 0) {
throw std::invalid_argument("invalid arg");
}

T* left = begin, *right = end;
T pivot = *left; //base
while (left < right) {
while (left < right && cmp(pivot, *right)) {
--right;
}
*left = *right;
while (left < right && cmp(*left, pivot)) {
++left;
}
*right = *left;
}
*left = pivot;

int length = left - begin + 1;
if (length > k) {
return quick_sort_topk(begin, left - 1, k, cmp);
} else if (length < k) {
return quick_sort_topk(right + 1, end, k-length, cmp);
}

return *left;
}

References: [1] 《BFPRT算法:时间复杂度O(n)求第k小的数字(分治算法+快排)》 [2] 《一文实现基本排序算法》 [3] 《拜托,面试别再问我TopK了》