如何基于随机选择快速排序获取数组的前K小元素？

我是算法初学者，知道可以用快速排序的随机选择算法获取数组的第K小元素，但现在需要修改这个算法，让它返回前K小的所有元素。比如数组int a[6] = { 1,3,4,2,5,6 }，调用random_select(a, 0, 5, 3)现在只会输出“3”，我期望输出“1 2 3”这类包含前3小的元素集合。

原代码 #

int random_select(int a[], int p, int r, int i) {

    if (p == r) {
        return a[p];
    }
    int q = partition_random(a, p, r);
    int k = q - p + 1;
    if (i == k) { //这就是我们要找的元素
        return a[q];
    }

    else if (i < k) { //在左侧部分
        return random_select(a, p, q - 1, i);
    }
    else { //右侧部分
        return random_select(a, q + 1, r, i - k);
    }
}

修改后的实现思路 #

原算法的核心是partition_random会把数组分成两部分：左边是小于等于基准的元素，右边是大于基准的元素。我们可以利用这个特性，不用只找单个元素，而是把所有属于前K小的元素都收集起来：

• 如果当前划分出的左半部分元素数量刚好等于K，那左半部分就是我们要的前K小元素；
• 如果左半部分元素数量比K多，说明前K小都在左半部分，递归处理左半部分；
• 如果左半部分元素数量比K少，那左半部分全部属于前K小，剩下的K-左半部分数量个元素去右半部分找，最后把两部分合并。

方式1：直接打印前K小元素 #

这种方式简单直接，适合快速验证结果：

void random_select_top_k(int a[], int p, int r, int k) {
    if (p > r || k <= 0) {
        return;
    }
    int q = partition_random(a, p, r);
    int left_size = q - p + 1;
    
    if (left_size == k) {
        // 输出当前左半部分的所有元素
        for (int i = p; i <= q; i++) {
            printf("%d ", a[i]);
        }
    } else if (left_size > k) {
        // 前K小都在左半部分，递归处理左半段
        random_select_top_k(a, p, q - 1, k);
    } else {
        // 先输出左半部分所有元素，再去右半部分找剩下的元素
        for (int i = p; i <= q; i++) {
            printf("%d ", a[i]);
        }
        random_select_top_k(a, q + 1, r, k - left_size);
    }
}

方式2：将结果存入指定数组 #

如果需要后续对前K小元素做进一步处理，可以把结果存入传入的数组中：

void random_select_top_k_store(int a[], int p, int r, int k, int result[], int *result_idx) {
    if (p > r || k <= 0) {
        return;
    }
    int q = partition_random(a, p, r);
    int left_size = q - p + 1;
    
    if (left_size == k) {
        // 把左半部分元素逐个存入结果数组
        for (int i = p; i <= q; i++) {
            result[(*result_idx)++] = a[i];
        }
    } else if (left_size > k) {
        // 递归左半部分找前K小
        random_select_top_k_store(a, p, q - 1, k, result, result_idx);
    } else {
        // 先存左半部分，再去右半部分找剩下的元素
        for (int i = p; i <= q; i++) {
            result[(*result_idx)++] = a[i];
        }
        random_select_top_k_store(a, q + 1, r, k - left_size, result, result_idx);
    }
}

使用示例 #

针对题目中的数组，调用方式如下：

// 方式1调用，直接打印结果
int a[6] = {1,3,4,2,5,6};
random_select_top_k(a, 0, 5, 3); // 输出：1 2 3 

// 方式2调用，将结果存入数组
int result[3];
int idx = 0;
random_select_top_k_store(a, 0, 5, 3, result, &idx);
// 此时result数组内容为{1,2,3}

注意事项 #

• 以上代码假设partition_random函数是把小于等于基准值的元素放到左半部分，如果你的划分逻辑是把大于等于基准的放左边，需要调整判断条件；
• 这个修改后的算法平均时间复杂度还是O(n)，比先排序再取前K个的O(n logn)更高效，适合处理大数据量的场景。

内容的提问来源于stack exchange，提问作者QuestionableCodingPractices