Day 084-Q3 | Randomized Quickselect + BFPRT 線形時間選択

問題

長さ $N$ の整数列 $A$ が与えられる。$Q$ 個のクエリ $(k_i)$ が与えられ、各クエリで $A$ の $k_i$ 番目に小さい要素（1-indexed）を答えよ。

パート1: 乱択クイックセレクト（Randomized Quickselect）を実装し、期待 $O(N)$ でクエリを処理せよ。

パート2: 決定論的線形時間選択アルゴリズム（Median-of-Medians / BFPRT）を実装し、最悪 $O(N)$ でクエリを処理せよ。

制約

パラメータ	範囲	備考
$N$	$\le 5 \times 10^5$	配列長
$Q$	$\le 10^5$	クエリ数
$k_i$	$1 \le k_i \le N$	1-indexed
$A_i$	$-10^9 \le A_i \le 10^9$	整数

入出力例

入力例1

8 5
3 1 4 1 5 9 2 6
1
4
7
8
3

出力例1

概念図: クイックセレクトと BFPRT の比較

ヒント

ヒント1（方向性）

「第 $k$ 小の要素を求める」問題は選択問題（Selection Problem）。ソートすれば $O(N \log N)$ だが、$O(N)$ で解ける。乱択クイックセレクトは期待 $O(N)$、BFPRT は最悪 $O(N)$（定数が大きい）。

ヒント2（アプローチ）

ヒント3（ほぼ答え）

import random

def quickselect(arr, k):  # k: 0-indexed
    while True:
        if len(arr) == 1:
            return arr[0]
        pivot = arr[random.randint(0, len(arr) - 1)]
        left  = [x for x in arr if x < pivot]
        mid   = [x for x in arr if x == pivot]
        right = [x for x in arr if x > pivot]
        if k < len(left):
            arr = left
        elif k < len(left) + len(mid):
            return pivot
        else:
            k -= len(left) + len(mid)
            arr = right

模範解答

import sys
import random
input = sys.stdin.readline

def solve():
    N, Q = map(int, input().split())
    A = list(map(int, input().split()))

    def quickselect(arr, k):
        """期待 O(N) 乱択クイックセレクト (0-indexed k)"""
        while True:
            if len(arr) == 1:
                return arr[0]
            pivot = arr[random.randint(0, len(arr) - 1)]
            left  = [x for x in arr if x < pivot]
            mid   = [x for x in arr if x == pivot]
            right = [x for x in arr if x > pivot]
            if k < len(left):
                arr = left
            elif k < len(left) + len(mid):
                return pivot
            else:
                k -= len(left) + len(mid)
                arr = right

    def median_of_medians(arr, k):
        """最悪 O(N) BFPRT アルゴリズム (0-indexed k)"""
        if len(arr) <= 5:
            return sorted(arr)[k]
        medians = []
        for i in range(0, len(arr), 5):
            chunk = sorted(arr[i:i+5])
            medians.append(chunk[len(chunk) // 2])
        pivot = median_of_medians(medians, len(medians) // 2)
        left  = [x for x in arr if x < pivot]
        mid   = [x for x in arr if x == pivot]
        right = [x for x in arr if x > pivot]
        if k < len(left):
            return median_of_medians(left, k)
        elif k < len(left) + len(mid):
            return pivot
        else:
            return median_of_medians(right, k - len(left) - len(mid))

    out = []
    for _ in range(Q):
        k = int(input()) - 1  # 0-indexed
        out.append(quickselect(A[:], k))

    sys.stdout.write('\n'.join(map(str, out)) + '\n')

solve()

Step-by-Step 解説

Step 1: クイックセレクトの期待計算量

毎回平均で半分に縮小するため $T(N) = T(N/2) + O(N) = O(N)$。最悪の場合は $O(N^2)$ だが、ランダムピボット選択により確率的に回避。

Step 2: BFPRT の最悪 $O(N)$ 保証

手順	内容	計算量
グループ化	5要素ずつ $N/5$ グループ	$O(N)$
中央値算出	各グループを定数時間でソート	$O(N)$
中央値の中央値	$N/5$ 要素に再帰	$T(N/5)$
分割	ピボットは 30-70% 保証	$O(N)$
再帰	最悪 $7N/10$ 要素	$T(7N/10)$

再帰式: $T(N) = T(N/5) + T(7N/10) + O(N) = O(N)$

Step 3: 実用上の選択

アルゴリズム	計算量	実用速度	用途
ソート	$O(N \log N)$	速い	複数クエリ
乱択 Quickselect	期待 $O(N)$	非常に速い	競プロ向き
BFPRT	最悪 $O(N)$	遅い（定数大）	理論的最適

よくあるミス

ミス	原因	正しい書き方
1-indexed / 0-indexed の混在	クエリ $k$ を変換し忘れる	`k -= 1` で 0-indexed に統一
pivot 選択で最悪ケース	常に最小や最大を選ぶ実装	`random.randint` でランダム選択
元の配列を破壊する	`arr` を直接変更	`A[:]` でコピーしてから渡す

次のステップ

発展問題: ストリーミング中央値（Dual Heap: Max-heap + Min-heap）
発展問題: ウェイトつき選択問題（重み付き中央値）

自己評価

理解度: 完全理解 / おおむね理解 / 要復習

自分の回答:

気づき・メモ: