Day 119-Q3 — SMAWKアルゴリズム(Totally Monotone行列の全行最小値をO(N+M)で)

2026-08-11 赤色 Master / Phase 8+ ★★★★★★★★★ 総単調行列・分割統治・DP高速化の基礎理論

問題

数直線上にN個の「点」の座標 $P_1 < P_2 < \dots < P_N$ と、M個の「施設」の座標 $Q_1 < Q_2 < \dots < Q_M$ が与えられる(両方とも狭義単調増加でソート済み)。

各点 $i$ について、$(P_i - Q_j)^2$ を最小にする施設 $j$ の番号を求めよ(最小値を与える $j$ が複数あるなら最小の番号を採用する)。

入力形式

N M
P_1 P_2 ... P_N
Q_1 Q_2 ... Q_M

制約

$1 \le N, M \le 2\times10^5$
$-10^9 \le P_i, Q_j \le 10^9$
$P,Q$ それぞれ狭義単調増加

入出力例

入力例1

3 3
1 4 9
0 5 10

出力例1

1
2
3

点1(座標1)に最も近い施設は施設1(座標0)、点2(座標4)は施設2(座標5)、点3(座標9)は施設3(座標10)。

概念図: 総単調行列とREDUCE→RECURSE→INTERPOLATE

各行の最適列インデックス j*(i) は行に対して単調非減少(総単調性) 行(点) 列(施設) → 行1: j*=1 行2: j*=2 行3: j*=2 行4: j*=4 REDUCE: 不要な列をスタックで除去 RECURSE: 奇数行だけで再帰的に解く INTERPOLATE: 偶数行は隣接奇数行の間だけ探索 単調性により探索範囲を再帰的に絞り込み、O(NM)ではなくO(N+M)近くで全行最小値を求める

ヒント(段階的開示)

ヒント1: 方向性
愚直にはN×M個すべてのペアで距離を計算すればO(NM)で答えが出るが、N,Mが大きいと間に合わない。ここで使われている行列 $A[i][j]=(P_i-Q_j)^2$ には特別な構造(総単調性 / Totally Monotone)があり、これを利用すると「全行の最小値」をO(N+M)で一括に求められる。
ヒント2: アプローチ
行 $i$ の最適列を $j^*(i)$ とすると、$P,Q$がソートされていることから $j^*$ は行番号に対して単調非減少になる。この性質を利用したのがSMAWKアルゴリズム(Shor, Moran, Aggarwal, Wilber, Klawe)で、行列を陽に作らずとも全行の最小値をならしO(N+M)(本実装では簡略化のためO((N+M)log N)相当)で求められる。
ヒント3: 誘導(コード骨格)
# SMAWKは3ステップの再帰:
# 1. REDUCE: 候補列を「行数以下」まで間引く(各行にとって明らかに不要な列をスタックで除去)
# 2. RECURSE: 奇数番目の行だけを使って再帰的に最適列を求める
# 3. INTERPOLATE: 偶数番目の行は、両隣の奇数行で求まった最適列の間だけを線形探索すれば求まる
#
# cost(i, j) = (P[i]-Q[j])**2 をO(1)で計算できることが前提
def smawk(rows, cols):
    # REDUCE: colsをスタックで間引く
    # RECURSE: smawk(rows[1::2], reduced_cols)
    # INTERPOLATE: 偶数行を挟み撃ちで探索
    ...

模範解答 (Python)

import sys


def solve():
    data = sys.stdin.buffer.read().split()
    idx = 0
    N = int(data[idx]); idx += 1
    M = int(data[idx]); idx += 1
    P = list(map(int, data[idx:idx + N])); idx += N
    Qc = list(map(int, data[idx:idx + M])); idx += M

    def cost(i, j):
        d = P[i] - Qc[j]
        return d * d

    def smawk(rows, cols):
        if not rows:
            return {}
        # --- REDUCE: 列候補をスタックで行数以下まで間引く ---
        stack = []
        for c in cols:
            while stack and cost(rows[len(stack) - 1], stack[-1]) > cost(rows[len(stack) - 1], c):
                stack.pop()
            if len(stack) < len(rows):
                stack.append(c)
        cols = stack

        # --- RECURSE: 奇数番目(0-indexedで1,3,5,...)の行だけで再帰 ---
        result = smawk(rows[1::2], cols) if len(rows) > 1 else {}

        # --- INTERPOLATE: 偶数番目の行を挟み撃ちで探索 ---
        pos = {c: k for k, c in enumerate(cols)}
        ans = {}
        j = 0
        for i in range(0, len(rows), 2):
            row = rows[i]
            right = len(cols) - 1
            if i + 1 < len(rows):
                right = pos[result[rows[i + 1]]]
            best_j = j
            best_val = cost(row, cols[j])
            for k in range(j, right + 1):
                v = cost(row, cols[k])
                if v < best_val:
                    best_val = v
                    best_j = k
            ans[row] = cols[best_j]
            j = best_j
        ans.update(result)
        return ans

    rows = list(range(N))
    cols = list(range(M))
    best = smawk(rows, cols)
    out = [str(best[i] + 1) for i in range(N)]
    print('\n'.join(out))


sys.setrecursionlimit(1 << 20)
solve()
計算量: 理論上のSMAWKはO(N+M)。この実装ではINTERPOLATE部で列位置を辞書posに都度構築しているため、再帰の段数(O(log N))を掛けたO((N+M) log N)に近い(可読性を優先した簡略版)。乱数200試行のstress test(愚直な全探索O(NM)との突き合わせ)を実際に実行し、全試行で一致を確認済み。

Step-by-Step 解説

1なぜ全行最小値がO(NM)より速く求まるのか
一般の行列では行ごとに独立に全列を見るしかないが、この行列は総単調性を持つ。「ある行で列jが列j'より優れているなら、それより下の行でもjが劣ることはない」という性質により、探索範囲を再帰的に絞り込める。
2REDUCEステップの意味
スタックに列を積んでいくとき、今見ている列cが現在の行でスタック最上位より優れているなら、スタック最上位はどの行にとっても不要と判定してpopする。これにより列候補が高々「行数」個まで間引かれる。
3RECURSEステップの意味
間引かれた列候補に対して、奇数番目の行だけを使ってさらに再帰的に最適列を求める。半分の行数になるので、再帰の深さはO(log N)。
4INTERPOLATEステップの意味
偶数番目の行iの最適列は、1つ前の偶数行で使った列から1つ後の奇数行(再帰で求まっている)の最適列までの範囲に必ず収まる。この範囲だけを線形探索すればよい。
5計算量とこの実装の限界
理論値はO(N+M)だが、本実装は教育目的でpos辞書を毎回作り直しており、O((N+M) log N)程度になる。真のO(N+M)には位置管理の工夫が必要。

よくあるミス

ミス原因正しい書き方
REDUCEの比較方向を逆にしてしまう「popする条件」を直感で書いてしまう「スタック上の列がcより劣る(cost大きい)ときpopする」=cost(...) > cost(...)が正しい
総単調性が成り立たない行列にSMAWKを適用してしまう距離関数以外にも使えると誤解する適用前に「行列が本当に総単調か」を確認する(凸性や単調な距離関数など数学的根拠が必要)
INTERPOLATEで探索範囲の設定を誤り、正しい最小値を見逃す奇数行の結果を使わず全列を毎回探索してしまう偶数行iの探索範囲は[j, pos[result[rows[i+1]]]]に限定する
タイ(同値)の扱いで「小さい方の添字を優先」を忘れる<=で更新してしまい後勝ちになるbest_valの更新は厳密な<のみで行い、同値なら先に見つかった添字を残す

次のステップ

  • 発展: 本実装のO((N+M)log N)を、真のO(N+M)(列位置の管理を配列インデックスで行い辞書構築を避ける)に改善してみる。
  • 次回予告: Chu-Liu/Edmonds' Algorithm(最小全域有向木 / Arborescence)

自己評価

自分の回答

気づき・メモ