【学习记录】线性选择

我们可以在 $O(n\log k)$ 的时间内获得一个序列内的第 $k$ 大的数。但是有没有更快的方法？

快速选择

一种方法是仿照快速排序，每次随机选择一个主元，并且将小于它和大于它的元素分别放到它的左边和右边。这样我们就很容易判断第 $k$ 大的元素在哪一部分，从而递归的去寻找。

最优情况下，我们每次找到的都是中位数，这样时间复杂度是线性的。但在极端情况下，对于某个序列和某个 $k$ ，该算法的时间复杂度会退化到 $O(n^2)$ 级别。而最原始的快速排序也有可能退化到这一复杂度。

BFPRT 算法

下面要介绍的 BFPRT 算法就是一个能在严格的线性时间内筛选出一个序列的第 $k$ 大的数的算法。它的精髓在于选取主元的过程。它并不是随机选取一个主元，而是反复调用自身进行主元的选取——也就是说，每次选取主元相当于就是在求解一个规模更小的线性选择问题。

它的步骤如下：

将 $n$ 个元素划分成为 $\lfloor \frac{n}{5}\rfloor$ 个组，每组 $5$ 个元素。其余的、长度不足 $5$ 的部分舍去。如果 $n<5$ 就自成一组。
用某种排序对每个组求中位数。这一步的时间复杂度是常数。
对于第 2 步中的 $\lfloor \frac{n}{5}\rfloor$ 个中位数调用 BFPRT 算法求出中位数，作为主元。

求出了主元之后，我们仍像快速选择一样，将小于它和大于它的元素分别放到它的左边和右边，并选取答案所在部分递归求解，直到获得答案。

这个算法的时间复杂度为 $O(n)$ ，具体证明可以参照《算法导论》。

这个算法的实现如下：

int BFPRT(int *a, int l, int r, int K);     // 前向声明；作用范围为 a[l, r]
int get_index(int *a, int l, int r){
    // 对 [l, r] 插入排序
    for (int i = l + 1; i <= r; ++i){
        int tmp = a[i], pos = i - 1;
        while (pos >= l && tmp < a[pos])
            a[pos + 1] = a[pos], --pos;
        a[pos + 1] = tmp;
    }
    return (l + r) >> 1;
}
int get_pivot(int *a, int l, int r){
    // 每五个获取中位数，不停的获取
    if (r - l < 4) 
        return get_index(a, l, r);
    int pos = l - 1; 
    for (int i = l; i + 4 <= r; i += 5) {
        int id = get_index(a, i, i + 4);
        swap(a[++pos], a[id]);
    }
    int len = pos - l + 1;
    return BFPRT(a, l, pos, (len >> 1) + 1);
}
int partition(int *a, int l, int r, int pivot_id){
    // 依照枢纽元划分
    swap(a[pivot_id], a[r]);
    int pos = l;
    for (int i = l; i < r; ++i)
        if (a[i] < a[r]){
            while (pos < i && a[pos] < a[r])
                ++pos;
            swap(a[pos], a[i]), ++pos;
        }
    swap(a[pos], a[r]);
    return pos;
}
int BFPRT(int *a, int l, int r, int K){
    // 主过程，和快速选择差不多
    int pivot_id = get_pivot(a, l, r);
    int div_pos = partition(a, l, r, pivot_id);
    if (div_pos - l >= K)
        return BFPRT(a, l, div_pos - 1, K);
    else if (div_pos - l + 1 == K)
        return div_pos;
    else 
        return BFPRT(a, div_pos + 1, r, K - div_pos + l - 1);
}

STL

事实上 STL 中已经包含了快速选择算法的实现，即 nth_element。它的使用用例如下：

int main(){
    // 从 0 开始计数，求第 i 小的数。
    int a[] = {-100, -90, -80, 1, 2, 3};
    int n; 
    cin >> n;           // n in [0, 6)
    nth_element(a, a + n, a + 6);
    cout << a[n] << endl;
    return 0;
}

参考资料

《算法导论（第 3 版）》