ABC012 D - バスと避けられない運命（ダイクストラ法解）

備忘録

問題

atcoder.jp

回答

import sys, os, math, bisect, itertools, collections, heapq, queue, copy, array

# from scipy.sparse.csgraph import csgraph_from_dense, floyd_warshall
# from decimal import Decimal
# from collections import defaultdict, deque

sys.setrecursionlimit(10000000)

ii = lambda: int(sys.stdin.buffer.readline().rstrip())
il = lambda: list(map(int, sys.stdin.buffer.readline().split()))
fl = lambda: list(map(float, sys.stdin.buffer.readline().split()))
iln = lambda n: [int(sys.stdin.buffer.readline().rstrip()) for _ in range(n)]

iss = lambda: sys.stdin.buffer.readline().decode().rstrip()
sl = lambda: list(map(str, sys.stdin.buffer.readline().decode().split()))
isn = lambda n: [sys.stdin.buffer.readline().decode().rstrip() for _ in range(n)]

lcm = lambda x, y: (x * y) // math.gcd(x, y)

MOD = 10 ** 9 + 7
INF = float('inf')


def dijkstra(N, start, graph):
    dist = [INF] * N  # 各頂点までの距離を保持する配列
    dist[start] = 0  # 開始始点の頂点は距離をゼロとする
    visited = [0] * N  # 各頂点の訪問状態を保持する配列[0: 未訪問, 1: 訪問済み]
    heap = [(0, start)]  # 距離と始点頂点（自分自身への移動距離はゼロ）
    heapq.heapify(heap)  # 距離と頂点の情報を距離で優先度付きキューとする
    while heap:
        # 頂点uとuに到達するまでの距離
        min_cost, u = heapq.heappop(heap)

        # さらに小さい距離で既に到達済みの場合はcontinue
        if dist[u] < min_cost: continue
        visited[u] = 1  # 頂点uを訪問済みにする

        # グラフの情報から接する頂点と移動距離を順に処理
        for cost, v in graph[u]:
            if visited[v] != 1:
                if dist[u] + cost < dist[v]:
                    # 未到達かつ小さい距離で移動できる場合には
                    # 距離の更新と優先度付きキューに(距離, 頂点)の情報を追加
                    dist[v] = dist[u] + cost
                    heapq.heappush(heap, (dist[v], v))
    return dist


def main():
    if os.getenv("LOCAL"):
        sys.stdin = open("input.txt", "r")

    N, M = il()
    graph = [[] for _ in range(N)]
    for _ in range(M):
        s, t, c = il()
        graph[s - 1].append((c, t - 1))
        graph[t - 1].append((c, s - 1))

    ret = INF
    for n in range(N):
        dist = dijkstra(N, n, graph)
        ret = min(ret, max(dist))
    print(ret)


if __name__ == '__main__':
    main()

考え方

以前、ワーシャルフロイド法で回答した問題をダイクストラ法で回答した。
AOJ ALDS1_12_C Single Source Shortest Path IIの回答が、この問題にも使えそうだなと思いついたため、とりあえず流してみたところ、何とかAC。
ちなみに、Python (3.8.2)で提出すると実行時間は4978 ms、
PyPy3 (7.3.0)で提出すると実行時間845msだった。