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

Posted 2020-11-16 shiina922

最直观的方法是先排序再取中位数, 时间复杂度 (O(nlog 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]) 为子列的元素个数.

[ egin{align*} mathbb E T(n) &le mathbb Eleft[Tleft(max(Y-1,n-Y) ight) + O(n) ight]& = sum_{k=1}^n frac1nmathbb Eleft[ Tleft(max(k-1,n-k) ight) ight] + O(n)& le frac2nsum_{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 ceil) 组, 每组 5 个, 最后一组可能不足 5 个. 用时 (O(n)).
2. 找到每组 5 个元素的中位数. 用时 (O(n)).
3. (递归地) 用 select 算法找到 (lceil n/5 ceil) 个中位数的中位数 (x). 用时 (T(lceil n/5 ceil)).
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) 大的元素个数至少有

[ 3left( leftlceil frac12 leftlceilfrac{n}{5} ight ceil ight ceil -2 ight) ge frac{3n}{10} - 6. ]

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

[ T(n) le T(lceil n/5 ceil) + 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.