Pythonで理解する蟻本「2-5 Layout（POJ No.3169）」(p.104) - クルトンのプログラミング教室

この記事は「プログラミングコンテストチャレンジブック第2版」（蟻本）の
「2-5 Layout（POJ No.3169）」(p.104)
のコードをPythonで書き直したものとなっています。

入力
入力例
解答

入力

$N\,ML\,MD\\AL_1\,BL_1\,DL_1\\:\\AL_{ML}\,BL_{ML}\,DL_{ML}\\AD_1\,BD_1\,DD_1\\:\\AD_{MD}\,BD_{MD}\,DD_{MD}$

入力例

4 2 1
1 3 10
2 4 20
2 3 3

解答

import sys

MAX_N = 1000
MAX_ML = 10000
MAX_MD = 10000
INF = float('inf')

# 入力
N, ML, MD = map(int,input().split())
AL = [0] * MAX_ML
BL = [0] * MAX_ML
DL = [0] * MAX_ML
for i in range(ML):
    AL[i], BL[i], DL[i] = map(int,input().split())
AD = [0] * MAX_MD
BD = [0] * MAX_MD
DD = [0] * MAX_MD
for i in range(MD):
    AD[i], BD[i], DD[i] = map(int,input().split())

d = [0] * MAX_N    # 最短距離
updated = False   # 更新されたか

def update(x, y, i):
    global updated
    if x > y:
        d[i] = y
        updated = True

# ベルマンフォード法によりdを計算する
def bellmanford():
    global updated
    for k in range(N + 1):
        updated = False
        # i+1からiへコスト0
        for i in range(N):
            if d[i + 1] < INF:
                update(d[i], d[i + 1], i)
        # ALからBLへコストDL
        for i in range(ML):
            if d[AL[i] - 1] < INF:
                update(d[BL[i] - 1], d[AL[i] - 1] + DL[i], BL[i] - 1)
        # BDからADへコスト-DD
        for i in range(MD):
            if d[BD[i] - 1] < INF:
                update(d[AD[i] - 1], d[BD[i] - 1] - DD[i], AD[i] - 1)

# 負閉路チェック
bellmanford()
if updated:
    print(-1)
    sys.exit()    # ここでプログラムを終了

d = [INF] * MAX_N
d[0] = 0
bellmanford()
res = d[N - 1]
if res == INF:
    res = -2
print(res)