D - 最安経路 / Cheapest Route Editorial by admin
Gemini 3.1 Pro (Thinking)概要
都市を頂点、道路を辺としたグラフにおいて、都市 \(1\) からいずれかの空港がある都市への最短経路(最小コスト)を求める問題です。
考察
都市間の移動にかかる通行料は、移動元の都市の人口と移動先の都市の人口の積で決まります。これは、各道路(辺)に対して固定のコスト \(M_a \times M_b\) があらかじめ設定されていると考えることができます。
また、各都市の人口 \(M_i\) は \(1\) 以上であるため、通行料(辺のコスト)は常に正の値になります。 このように、辺のコストが非負であるグラフの単一始点最短経路問題は、ダイクストラ法(Dijkstra’s Algorithm) を用いて効率的に解くことができます。
もし、辺のコストを考慮しない単純な幅優先探索(BFS)を使ってしまうと正しい最小コストが求まりません。また、ベルマンフォード法を使うと計算量が \(O(N \times K)\) となり、今回の制約(\(N, K \le 2 \times 10^5\))では実行時間制限(TLE)に引っかかってしまいます。そのため、優先度付きキュー(ヒープ)を用いたダイクストラ法を採用するのが最適です。
アルゴリズム
- グラフの構築: 各道路が結ぶ都市 \(U_k\) と \(V_k\) について、コストを \(M_{U_k} \times M_{V_k}\) とした無向辺をグラフに追加します。 (具体例:都市 \(2\) の人口が \(3\)、都市 \(3\) の人口が \(5\) で、これらを結ぶ道路がある場合、この辺のコストは \(3 \times 5 = 15\) になります。)
- 最短距離の計算(ダイクストラ法):
- 始点である都市 \(1\) の距離を \(0\)、それ以外の都市の距離を無限大(\(\infty\))で初期化します。
- 優先度付きキューに
(距離 0, 都市 1)を追加します。 - キューから現在最も距離が短い都市を取り出し、その都市から繋がる隣接都市へ移動したときの距離を計算します。それが現在記録されている距離よりも短ければ値を更新し、キューに新しい距離と都市のペアを追加します。
- これをキューが空になるまで繰り返します。
- 答えの導出: 空港のある都市 \(E_1, E_2, \dots, E_P\) について、計算した都市 \(1\) からの最短距離を確認し、その中で最も小さい値を答えとして出力します。
計算量
- 時間計算量: \(O((N + K) \log N)\) グラフの構築に \(O(N + K)\)、ダイクストラ法による最短経路の計算に \(O((N + K) \log N)\) かかります。制約の最大値でも実行時間制限に余裕で間に合います。
- 空間計算量: \(O(N + K)\) グラフの隣接リストの保持に \(O(N + K)\)、距離配列や優先度付きキューに \(O(N)\) のメモリを使用します。
実装のポイント
高速な入出力: Pythonの場合、入力データが非常に大きくなる可能性があるため、
sys.stdin.read().split()を使って入力を一括で読み込むと実行時間を大幅に短縮できます。無限大の初期値: 距離の初期値(無限大)は、答えが取り得る最大値よりも十分に大きい値(例えば \(10^{18}\))に設定する必要があります。
古い情報のスキップ: 優先度付きキューから取り出した要素について、
if d > dist[u]: continueのようにすでに最短距離が更新されている古い情報をスキップすることで、無駄な探索を省くことができます。ソースコード
import sys
import heapq
def main():
input_data = sys.stdin.read().split()
if not input_data:
return
N = int(input_data[0])
K = int(input_data[1])
P = int(input_data[2])
M = [0] * (N + 1)
for i in range(1, N + 1):
M[i] = int(input_data[2 + i])
adj = [[] for _ in range(N + 1)]
idx = 3 + N
for _ in range(K):
u = int(input_data[idx])
v = int(input_data[idx + 1])
idx += 2
cost = M[u] * M[v]
adj[u].append((v, cost))
adj[v].append((u, cost))
INF = 10**18
dist = [INF] * (N + 1)
dist[1] = 0
pq = [(0, 1)]
while pq:
d, u = heapq.heappop(pq)
if d > dist[u]:
continue
for v, cost in adj[u]:
nd = d + cost
if dist[v] > nd:
dist[v] = nd
heapq.heappush(pq, (nd, v))
ans = min(dist[int(input_data[i])] for i in range(idx, idx + P))
print(ans)
if __name__ == '__main__':
main()
この解説は gemini-3.1-pro-thinking によって生成されました。
posted:
last update: