ROSALIND｜Median (MED)

發表於 更新於

給定長度為 n 的整數數列，求其中第 k 小的數字。

Given: A positive integer $n \leq 10^5$ and an array $A[1..n]$ of integers from $-10^5$ to $10^5$, and a positive number $k \leq n$.

Return: The $k$-th smallest element of $A$.

(https://rosalind.info/problems/med/)

我想在介紹解題方法前先釐清題目的定義。假設數列由 1 開始計算（這是 Rosalind 題目的慣例了），所謂第 k 小的數字其實是遞增排序後，數列由左至右數來第 k 個數字。因此，數列 [4, 4, 5, 6] 第 2 小的數字是 4 而非 5，這可能跟想像不太一樣。

前述定義也相當於解題方法，它的時間複雜度取決於採用的排序算法，Python 的 sorted 採用 Timsort（參考官方文件，最低可達 $O(nlogn)$。

1
2
3
def med_sort(l: list[int], k: int) -> int:
    """Get the kth (1-based) smallest integer within a list (sort method)"""
    return sorted(l)[k-1] # 因為是 1-based 所以要減 1

另一種方法則利用了 heap 的特性，如果從數列取出前 k 小的數字建立 max heap，它的根節點即為數列第 k 小的數字。實作方法可參考我之前寫 partial sort 的介紹，不過 Python 裡面有 heapq 套件可以達到相同目的。

因為不用處理調整整個數列，所以時間複雜度略小於排序法，約為 $O(nlogk)$

1
2
3
4
5
import heapq
def med_hp(l: list[int], k: int) -> int:
    """Get the kth (1-based) smallest integer within a list (heap method)"""
    smallest = heapq.nsmallest(k, l)
    return smallest[-1]

最後一種是稱為 quick select 的算法，它從數列隨機選取一個基準值，將數列分為小於、等於和大於基準值等部分。接著判斷 k 位於那個部分，就針對那個區域重複劃分，直到 k 落在中間區域為止。

這個算法的時間複雜度取決於數列的數字分布和基準值的選取，如果碰到一個遞減排序的數列，又好死不死選取第一個數字當基準值，意味著要對整個數列每個位置進行分割。

為了減少這種狀況，通常會隨機選取基準值，又或是採用一個很玄的算法，叫 median of medians，它的功能是透過更有效率的方式選取基準值，避免落入最差的情境讓整個算法的平均時間降低到 O(n)。

我讀了這算法的內容，覺得玄之又玄（也就是說我還沒讀懂），它將數列五個一組分段，然後選取各自中位數的中位數當作基準值，藉此避免選到過於極端的值（偏左或偏右）。之所以感到訝異是因為，分組的尺寸既不是三也不是七，而是選擇五。在給人十分嚴謹的算法裡面居然出現像是經驗法則的參數決定依據，讓我覺得十分有趣。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
def partition(l: list[int], pivot: int) -> list[int]:
    """3-way partition"""
    left, mid, right = [], [], []
    for x in l:
        if x < pivot:
            left.append(x)
        elif x > pivot:
            right.append(x)
        else:
            mid.append(x)
    return left, mid, right
def med_qs(l: list[int], k: int) -> int:
    """Get the kth (1-based) smallest integer within a list (quick select method)"""
    if len(l) == 1:
        return l[0]
    pivot = random.choice(l)
    left, mid, right = partition(l, pivot)
    if k <= len(left): # k locate at left part
        return med_qs(left, k)
    elif k > len(left) + len(mid): # k locate at right part
        return med_qs(right, k - len(left) - len(mid)) # 1-based, do not need to minus 1
    else:
        return mid[0]