如何使用分区解决重复的top-k问题？

Question

给定一个数组和一个数字 k，其中 k 小于数组的大小，我们需要找到给定数组中第 k 个最小的元素。

当数组中有重复项时，我很困惑。 显然，当试图在 [2, 2, 2, 2, 2] 中找到第三个元素时，我的下面的算法将被卡住 如何处理这样的问题？

下面的函数使用快速选择算法找出第k个最小的元素。

int findTopK(int L, int R, int K) {
    if (L == R) return arr[L];
    int i = L, j = R;
    int pivot = arr[R];
    while (true) {
        while (i <= R && arr[i] < pivot) i++;
        while (j >= L && arr[j] >= pivot) j--;
        if (i >= j) break;
        swap(arr[i], arr[j]);
        i++; j--;
    }
    //[L, j] [j+1, R]
    //left: < pivot
    //right: >= pivot

    if (j - L + 1 < K) return findTopK(j+1, R, K - (j - L + 1));
    else return findTopK(L, j, K);
}

Answer 1

最后我找到了一个避免陷入上述情况的解决方案。 该功能可以不断缩小目标区间的范围。 此外，@BeyelerStudios，fat pivot 正是这个问题的标准解决方案。感谢您的慷慨分享！

int findTopK(int L, int R, int K) {
    int i = L, j = R - 1;
    int pivot = arr[R];
    while (true) {
        while (i <= R && arr[i] < pivot) i++;
        while (j >= L && arr[j] >= pivot) j--;
        if (i >= j) break;
        swap(arr[i], arr[j]);
    }
    swap(arr[R], arr[j+1]);
    //[L, j], j+1, [j+2, R]
    //left: < pivot
    //equal: pivot
    //right: >= pivot

    if (j+1 - L + 1 == K) return pivot;

    if (j+1 - L + 1 < K) return findTopK(j+2, R, K - (j+1 - L + 1));
    else return findTopK(L, j, K);
}

Answer 2

基于“经典”Hoare 分区方案的替代方法：

int QuickSelectr(int a[], int lo, int hi, int k )
{
    if (lo == hi)                           // recurse until lo == hi
        return a[lo];
    int p = a[(lo+hi)/2];                   // Hoare partition
    int i = lo - 1;
    int j = hi + 1;
    while (1){
        while (a[++i] < p);
        while (a[--j] > p);
        if (i >= j)
            break;
        std::swap(a[i], a[j]);
    }
    if(k <= j)
        return QuickSelectr(a, lo, j-0, k); // include a[j]
    else
        return QuickSelectr(a, j+1, hi, k); // exclude a[j]
}

// parameter check
int QuickSelect(int *a, int lo, int hi, int k)
{
    if(a == (int *)0 || k < lo || k > hi || lo > hi)
        return 0;
    return QuickSelectr(a, lo, hi, k);
}

---

示例测试代码。在我的系统上，N == 8 大约需要 5 秒，N == 9 大约需要 2 分钟。

#define N (8)
int main()
{
int a[N] = {0};
int b[N];
int c[N];
int i, j;
    while(1){
        memcpy(b, a, N*sizeof(int));
        std::sort(b+0, b+N);
        for(j = 0; j < N; j++){
            memcpy(c, a, N*sizeof(int));
            i = QuickSelect(c, 0, N-1, j);
            if(i != b[j])
                break;
        }
        i = N-1;
        while(1){
            if(++a[i] != N || i == 0)
                break;
            a[i--] = 0;
        }
        if(a[0] == N)
            break;
    }
    if(a[0] == N)
        printf("passed\n");
    else
        printf("failed\n");         
    return 0;
}