ABC191 E - Come Back Quickly - ペンギンかっこいい

備忘録

問題

回答

import sys
import os
import math
import bisect
import itertools
import collections
import heapq
import queue
import array
import time
import datetime

# 時々使う
# import numpy as np
# from decimal import Decimal, ROUND_HALF_UP
# from scipy.sparse.csgraph import csgraph_from_dense, floyd_warshall

# 再帰の制限設定
sys.setrecursionlimit(10000000)


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


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


def lcm(x, y): return (x * y) // math.gcd(x, y)


MOD = 10 ** 9 + 7
# MOD = 998244353
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)]
    dist = [INF] * N
    for _ in range(M):
        a, b, c = il()
        graph[a - 1].append((c, b - 1))
        if a == b:
            dist[a-1] = min(dist[a-1], c)

    di = []
    for i in range(N):
        di.append(dijkstra(N, i, graph))

    for i in range(N):
        for j in range(N):
            if i == j:
                continue
            dist[i] = min(di[j][i] + di[i][j], dist[i])

    for d in dist:
        if d == INF:
            print(-1)
        else:
            print(d)


if __name__ == '__main__':
    main()

考え方

上記のコードはPyPy3 (7.3.0)で提出時には実行時間 896msでした。
Python (3.8.2)で提出した場合にはTLE（実行時間制限超過）となります。

町を頂点、道路を辺と置き換えてグラフ問題として考えます。
つまり、ある頂点から１つ以上の辺を通り元の頂点に戻るためのコストを考える問題です。
また、辺は一方通行なため有向グラフとして扱います。

回答はダイクストラ法を用いて回答しました。
以前に作った優先度キューを用いたダイクストラで
頂点Aから異なる頂点Bへの最短距離を求め、頂点Bから頂点Aまでの最短距離の合計を正しい散歩経路とします。
自分自身を参照する辺が存在する場合には、
他の頂点を経由した場合と自分自身を参照する頂点とでコストの低い方を回答とします。