B - 最寄りの避難所 / Nearest Shelter Editorial by admin
gemini-3-flash-thinking概要
数直線上に並んだ \(N\) 軒の家それぞれについて、最も近い場所にある避難所(計 \(M\) 箇所)との距離を求める問題です。
考察
素朴な方法
各家について、すべての避難所との距離を計算して最小値を探す方法が考えられます。しかし、この方法では \(1\) 軒あたり \(M\) 回の比較が必要になり、全体で \(O(N \times M)\) の計算量となります。 本問題では \(N, M \leq 2 \times 10^5\) であるため、最大で \(4 \times 10^{10}\) 回程度の計算が必要になり、実行時間制限に間に合いません(TLEとなります)。
効率的な方法
避難所の座標 \(P_1, P_2, \ldots, P_M\) が昇順に並んでいることに注目します。 ある家 \(X_i\) にとって最も近い避難所は、「\(X_i\) 以上の座標にある避難所のうち最も左にあるもの」か、「\(X_i\) 未満の座標にある避難所のうち最も右にあるもの」のいずれかです。
例えば、避難所が座標 \(10, 40, 70\) にあり、家が座標 \(50\) にある場合、近い候補は \(40\) か \(70\) のどちらかになります。
このように「ソート済みの配列から特定の値を挿入すべき位置」を探すには、二分探索(Binary Search)を用いることで、1軒あたり \(O(\log M)\) で候補を見つけることができます。
アルゴリズム
- 各家 \(X_i\) について、以下の手順を繰り返します。
- 二分探索(Pythonでは
bisect_left関数)を用いて、避難所のリスト \(P\) の中で \(X_i\) 以上の値が初めて現れるインデックスidxを探します。 idxの値に応じて、以下の 3 つのケースに分けて距離を計算します。idx == 0の場合: すべての避難所が家より右側にあります。最も近いのは \(P[0]\) です。idx == Mの場合: すべての避難所が家より左側にあります。最も近いのは \(P[M-1]\) です。- それ以外の場合: 家のすぐ右側に \(P[idx]\)、すぐ左側に \(P[idx-1]\) という 2 つの候補があります。両方の距離を計算し、小さい方を採用します。
- 全ての家の計算結果を順に出力します。
計算量
- 時間計算量: \(O(N \log M)\)
- 各家(\(N\) 軒)に対して、避難所のリスト(サイズ \(M\))を二分探索するためです。
- 空間計算量: \(O(N + M)\)
- 入力された家と避難所の座標を保持するためのメモリが必要です。
実装のポイント
二分探索の活用: Python の標準ライブラリ
bisectモジュールにあるbisect_leftを使うと、簡潔かつ高速に探索が行えます。高速な入出力: \(N, M\) が大きいため、
input()ではなくsys.stdin.readを使ったり、結果をリストにまとめて最後に一度にjoinして出力したりすることで、入出力によるタイムロスを減らしています。境界条件の処理: 家が避難所の範囲よりも外側(左端よりも左、または右端よりも右)にある場合に、配列のインデックス外を参照しないよう注意が必要です。
ソースコード
import sys
from bisect import bisect_left
def main():
input_data = sys.stdin.read().split()
if not input_data:
return
N = int(input_data[0])
M = int(input_data[1])
X = list(map(int, input_data[2 : 2 + N]))
P = list(map(int, input_data[2 + N : 2 + N + M]))
results = []
for x in X:
idx = bisect_left(P, x)
if idx == 0:
results.append(P[0] - x)
elif idx == M:
results.append(x - P[M-1])
else:
dist_right = P[idx] - x
dist_left = x - P[idx-1]
if dist_right < dist_left:
results.append(dist_right)
else:
results.append(dist_left)
sys.stdout.write('\n'.join(map(str, results)) + '\n')
if __name__ == '__main__':
main()
この解説は gemini-3-flash-thinking によって生成されました。
posted:
last update: