$O(n)$ 时间找到中位数

最直观的方法是先排序再取中位数, 时间复杂度 \(O(n\log n)\). 然而最近才得知中位数有时间复杂度 \(O(n)\) 的算法, 事实上任意顺序统计量都可以用 \(O(n)\) 时间找出.

In Expected Linear Time

记待排序的数列 \(A = [a_0, \dots, a_{n-1}]\), 其子列 \(A[p, \dots, r] = [a_p, \dots, a_r]\). 类似快排的方法, 找第 \(i\) 小的元素.

def partition(A, p, r):
    pivot = A[r]
    i = p - 1
    for j in range(p, r):
        if A[j] <= pivot:
            i += 1
            A[i], A[j] = A[j], A[i]
    i += 1
    A[i], A[r] = A[r], A[i]
    return i
    
def randomized_partition(A, p, r):
    i = random.randint(p, r)
    A[i], A[r] = A[r], A[i]
    return partition(A, p, r)
    
def randomized_select(A, p, r, i):
    if p == r:
        return A[p]
    q = randomized_partition(A, p, r)
    k = q - p + 1
    if i == k:  
        # the pivot value is the answer
        return A[q]
    elif i < k:
        return randomized_select(A, p, q-1, i)
    else:
        return randomized_select(A, q+1, r, i-k)

记时间复杂度为 \(T(n)\), 以及 \(Y = \#A[p, \dots, r]\) 为子列的元素个数.

\[\begin{align*} \mathbb E T(n) &\le \mathbb E\left[T\left(\max(Y-1,n-Y)\right) + O(n)\right]\\ & = \sum_{k=1}^n \frac1n\mathbb E\left[ T\left(\max(k-1,n-k)\right) \right] + O(n)\\ & \le \frac2n\sum_{k=[n/2]}^{n-1}\mathbb ET(k) + O(n). \end{align*} \]

之后易证 (substitution method) \(\mathbb ET(n) = O(n)\). 不过 worst-case 是 \(O(n^2)\).

In Worst-Case Linear Time

依然是从 \(n\) 个元素的数列中找第 \(i\) 小的元素. 不妨约定, 当偶数个元素时, 中位数取中间两个数中较小的那个.

算法记为 select 算法, 总体和前一个算法一样, 关键是找到一个好的 pivot. 记时间复杂度为 \(T(n)\).

1. 把数列分成 \(\lceil n/5\rceil\) 组, 每组 5 个, 最后一组可能不足 5 个. 用时 \(O(n)\).
2. 找到每组 5 个元素的中位数. 用时 \(O(n)\).
3. (递归地) 用 select 算法找到 \(\lceil n/5\rceil\) 个中位数的中位数 \(x\). 用时 \(T(\lceil n/5\rceil)\).
4. 以 \(x\) 为 pivot 进行划分 (partition), 记 \(x\) 为第 \(k\) 小的元素. 用时 \(O(n)\).
5. 若 \(i = k\), 返回 \(x\); 若 \(i < k\), (递归地) 用 select 算法在较小的 partition 中找第 \(i\) 小的元素; 若 \(i > k\), 则在较大的 partition 中找第 \(i - k\) 小的元素.

下面考虑第 5 步的用时. 考虑比 \(x\) 大的元素个数, 有一半的组, 每组 3 个元素比 \(x\) 大 (除了 \(x\) 所在的组和最后一个不满 5 个元素的组以外). 故比 \(x\) 大的元素个数至少有

\[3\left( \left\lceil \frac12 \left\lceil\frac{n}{5}\right\rceil \right\rceil -2 \right) \ge \frac{3n}{10} - 6. \]

故第 5 步最多用时 \(T(7n/10 + 6)\). 因此

\[T(n) \le T(\lceil n/5\rceil) + T(7n/10 + 6) + O(n). \]

易证 \(T(n) = O(n)\).

注意到若分为每组 3 个, 则不能如上证明 \(T(n) = O(n)\).

References

Leiserson, C. E., Rivest, R. L., Cormen, T. H., & Stein, C. (2009). Introduction to algorithms (3rd ed.) (pp. 215-222). Cambridge, MA: MIT press.