[Python] ABC012 D ダイクストラ法

問題

D – バスと避けられない運命

回答

ダイクストラ法を使います。

Python では TLE でしたが、PyPy では AC でした。

import collections
import sys
import heapq

N, M = map(int, input().split())
edges = [list(map(int, input().split())) for _ in range(M)]

# 隣接リスト
adjacent_dict = collections.defaultdict(list)
for frm, to, distance in edges:
    adjacent_dict[frm].append([to, distance])
    adjacent_dict[to].append([frm, distance])


ans_dict = collections.defaultdict(lambda: sys.maxsize)

# 全てのスタート地点に対してダイクストラ法を適用
for start in range(1, N+1):
    
    # 全ての距離を最大値に指定
    distance_dict = collections.defaultdict(lambda: sys.maxsize)
    # スタート地点の距離は0
    distance_dict[start] = 0

    # (距離, 頂点) の優先度付きキュー
    priority_queue = []
    heapq.heappush(priority_queue, (0, start))

    # 訪問済みフラグ
    visited = collections.defaultdict(bool)

    while len(priority_queue):
        # 距離が最小であるノードをを取得
        cur_distance, cur_node = heapq.heappop(priority_queue)
        # 訪問済みフラグ
        if visited[cur_node]:
                continue
        visited[cur_node] = True

        # 隣接するノードとの距離    
        for next_node, distance in adjacent_dict[cur_node]:
            if visited[next_node]:
                continue
            new_distance = cur_distance + distance
            # 新しい距離が短い場合は距離を更新
            if distance_dict[next_node] > new_distance:
                distance_dict[next_node] = new_distance

                heapq.heappush(priority_queue, (new_distance, next_node))
    
    # そのスタート地点で取り得る最大の距離
    ans_dict[start] = max(distance_dict.values())

# 全てのスタート地点の中で最小の距離
print(min(ans_dict.values()))