2008年12月1日星期一

topK

设U是一个全序集,A是U上的一个数组,长度为n。topK算法把最大的K个元素放在数组A的最前面。
思路类似快速排序,每次划分后只递归topK一边,而不是两边:
划分:在A中随便取一个pivot放中间,把大于等于pivot的数放左边,小于pivot的数放右边。设划分后pivot的下标为p。
如果p=K,太好了,左边的K个数刚好就是最大的。
如果p<K,我们还需要更多的比较大的数,从右侧的数中跑topK-p。
如果p>K,说明我们一下找多了,需要精益求精,在左侧的数中跑topK。
用C++写的,注意其中需要T重载一些比较运算符:
template<typename T>
void top(T *a, size_t n, size_t k) {
    if (n <= 0 || k <= 0 || n <= k) {
        printf("top need n>k! n=%d k=%d\n", n, k);
        throw "";
    }
    while (true) {
        int j = rand() % n;
        T pv = a[j];
        swap(a, 0, j);
        size_t p = 1;
        for (size_t i = 1; i < n; i++) {
            if (a[i] >= pv) {
                swap(a, i, p);
                p++;
            }
        }
        j = p - k;
        if (j < 0) {
            a = a + p;
            n = n - p;
            k = k - p;
        } else if (j > 0) {
            while (a[--p] == pv && p>k)
                ;
            swap(a, 0, p);
            if (p == k)
                break;
            else
                n = p;
        } else
            break;
    }
}
边界条件的处理略微有些复杂,我搞了好久才搞好。为了保证循环有限步完成,必须要让问题的规模确实减小。
为了加速处理所有数都相同的输入,p>K时,在不破坏正确性的、且不增加复杂度的情况下p尽量取小些。

没有评论: