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