Day 133-Q5 | BFPRT法／Median of Medians（決定的選択アルゴリズム）

問題

$N$個の整数からなる数列$A$（重複あり得る）と整数$K$（$1\le K\le N$）が与えられる。組み込みのソート関数・heapq.nlargest/nsmallestを使わず、最悪計算量$O(N)$が保証される決定的選択アルゴリズム（Median of Medians / BFPRT）を実装して、$A$を昇順に並べたときの$K$番目の値を出力せよ。

入力形式

N K
a_1 a_2 ... a_N

制約

$1 \le N \le 2\times10^5$

$1 \le K \le N$

$-10^9 \le a_i \le 10^9$

入出力例

入力例1

9 4
7 2 9 4 1 8 3 6 5

出力例1

入力例2

5 1
5 5 5 5 5

出力例2

概念図

ヒント（段階的開示）

ヒント1（方向性）

クイックセレクトの弱点はピボット選びが偏ると最悪$O(N^2)$になること。もしどんな入力でも「配列全体の一定割合を両側に残せる」ピボットを必ず$O(N)$で見つけられるなら、再帰は$T(N)=T(cN)+O(N)$（$c<1$）の形になり全体$O(N)$になる。そのようなピボットをどう作るか？

ヒント2（アプローチ）

配列を5個ずつのグループに分け、各グループの中央値を求める（$O(N)$）。その中央値の集合に対して再帰的に同じ選択アルゴリズムを適用し「中央値たちの中央値」をピボットとして使う。組み合わせ論から、配列全体の約30%以上を「ピボットより小さい」「ピボットより大きい」のそれぞれに保証できる。

ヒント3（誘導）

def select(lst, i):
    n = len(lst)
    if n <= 5:
        return sorted(lst)[i]
    medians = []
    for g in range(0, n, 5):
        group = lst[g:g+5]
        medians.append(sorted(group)[len(group)//2])
    mom = select(medians, (len(medians)-1)//2)
    lows   = [x for x in lst if x < mom]
    pivots = [x for x in lst if x == mom]
    highs  = [x for x in lst if x > mom]
    if i < len(lows):
        return select(lows, i)
    elif i < len(lows)+len(pivots):
        return mom
    else:
        return select(highs, i - len(lows) - len(pivots))

模範解答 (Python)

import sys
sys.setrecursionlimit(10000)


def select(lst, i):
    n = len(lst)
    if n <= 5:
        return sorted(lst)[i]

    medians = []
    for g in range(0, n, 5):
        group = lst[g:g + 5]
        group.sort()
        medians.append(group[len(group) // 2])

    mom = select(medians, (len(medians) - 1) // 2)

    lows = []
    pivots = []
    highs = []
    for x in lst:
        if x < mom:
            lows.append(x)
        elif x > mom:
            highs.append(x)
        else:
            pivots.append(x)

    if i < len(lows):
        return select(lows, i)
    elif i < len(lows) + len(pivots):
        return mom
    else:
        return select(highs, i - len(lows) - len(pivots))


def solve():
    data = sys.stdin.buffer.read().split()
    n = int(data[0]); k = int(data[1])
    a = [int(x) for x in data[2:2 + n]]
    print(select(a, k - 1))

solve()

注: Pythonのsorted()はC実装のTimsortで実測上非常に高速なため、実務ではsorted(a)[k-1]で十分なことも多い。本問の狙いは乱数に頼らず最悪計算量O(N)を理論的に保証するBFPRTの構成そのものを理解すること。

Step-by-Step 解説

15人組に分けて中央値を集める
配列を5個ずつのグループに分割し各グループをソートして中央値を取り出す。全体O(N)。

2中央値の中央値を再帰的に求める
集めた中央値のリスト（サイズN/5）に対して同じselect関数を再帰適用しmomを求める。

3なぜ良いピボットになるか
momを含むグループの半分では中央値以上の要素が3個確保されるため、momより大きい要素はおよそ3N/10個以上保証される（対称に小さい側も同様）。

4漸化式を解く
$T(N)=T(N/5)+T(7N/10)+O(N)$。$1/5+7/10=9/10<1$なので各レベルの作業量が幾何級数的に減少しO(N)に収束。

5クイックセレクトとの違い
乱択クイックセレクトは期待値O(N)だが最悪O(N^2)。BFPRTは入力に依存せず最悪計算量O(N)を保証する。

よくあるミス

ミス	原因	正しい書き方
グループサイズを5以外にしてしまう	「小さいグループなら何でもいい」と誤解	5が理論上の最小の実用値。3だと係数の和が1未満にならない
中央値集合への再帰を忘れ普通のソートで済ませる	median of mediansをもう一度selectせず妥協	select(medians,...)を再帰呼び出しし正確な中央値を求める
3-way partitionをせず重複要素で無限ループ	全要素が同じ値のとき2分割では停止しない場合がある	pivots(momと等しい要素)を明示的に分離しiがその範囲なら即momを返す
highsへの再帰でインデックス補正を忘れる	lowsとpivotsの分だけ順位がずれることを見落とす	select(highs, i-len(lows)-len(pivots))とオフセットを引く

次のステップ

発展: リスト内包表記による3-way partitionはO(N)の追加メモリを使う。in-place partitionに書き換え空間計算量をO(log N)に落とせ。
発展2: BFPRTはKirkpatrick–Seidelの出力鋭敏凸包アルゴリズムなど他の分割統治アルゴリズムの内部部品としても使われる。どこで使われうるか調べてみよ。
次回予告: 次回のテーマは今後のセッションで選定

自己評価

完全理解おおむね理解要復習

Day 133-Q5 — BFPRT法／Median of Medians（最悪計算量O(N)の決定的選択アルゴリズム）

問題