如何修改快速排序实现平均O(n)时间复杂度查找数组中位数

实现思路 #

你要的这个算法叫做快速选择，核心是复用快排的分区逻辑，每次只递归中位数所在的半区，不需要对整个数组排序，平均时间复杂度可以达到O(n)。
首先明确中位数的目标位置：设数组长度为n，我们要找的是排序后索引为k = n / 2的元素（索引从0开始，奇数长度取中间元素，偶数长度默认取下中位数，要上中位数的话改k为(n-1)/2即可）。
每次调用partition后，基准值pivot会落在最终排序的正确位置pi上：

• 若pi == k：直接返回arr[pi]就是中位数
• 若pi > k：中位数在左半区间[low, pi-1]，仅递归左半区
• 若pi < k：中位数在右半区间[pi+1, high]，仅递归右半区

完整修改后代码 #

原来的swap和partition函数完全不用改动，只需要替换原来的quickSort函数，新增快速选择和中位数查找入口即可：

// 原有swap、partition函数保持不变
static void swap(int[] arr, int i, int j){
    int temp = arr[i];
    arr[i] = arr[j];
    arr[j] = temp;
}

static int partition(int[] arr, int low, int high){
    int pivot = arr[high]; 
    int i = (low - 1); 

    for(int j = low; j <= high - 1; j++){
        if (arr[j] < pivot) {
            i++; 
            swap(arr, i, j);
        }
    }
    swap(arr, i + 1, high);
    return (i + 1);
}

// 新增快速选择函数
static int quickSelect(int[] arr, int low, int high, int k) {
    int pi = partition(arr, low, high);
    if (pi == k) {
        return arr[pi];
    } else if (pi > k) {
        // 只递归左半区
        return quickSelect(arr, low, pi - 1, k);
    } else {
        // 只递归右半区
        return quickSelect(arr, pi + 1, high, k);
    }
}

// 中位数查找入口
static double findMedian(int[] arr) {
    int n = arr.length;
    // 奇数长度直接取中间值
    if (n % 2 == 1) {
        return quickSelect(arr, 0, n - 1, n / 2);
    } else {
        // 偶数长度取中间两个数的平均值，若业务允许只取单个中位数，直接返回quickSelect(arr, 0, n-1, n/2)即可
        int leftMid = quickSelect(arr, 0, n - 1, (n / 2) - 1);
        int rightMid = quickSelect(arr, 0, n - 1, n / 2);
        return (leftMid + rightMid) / 2.0;
    }
}

优化提示 #

如果要避免最坏时间复杂度O(n²)，可以在每次分区前随机交换区间内的一个元素和区间末尾的元素，把pivot改成随机选的，这样最坏情况出现的概率几乎为0。

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