クルトンのプログラミング教室

Pythonの使い方やPythonを使った競技プログラミングの解法などを解説しています。

Pythonで理解する蟻本「2-5 単一始点最短路問題1(ベルマンフォード法)」(p.95)

この記事は「プログラミングコンテストチャレンジブック第2版」(蟻本)の
「2-5 単一始点最短路問題1(ベルマンフォード法)」(p.95)
のコードをPythonで書き直したものとなっています。

入力

蟻本には負の閉路がないグラフだけが入力例として載っていたので、負の閉路があるグラフも入力例として載せておきます。

また、蟻本に載っているグラフは無向グラフですが、このブログで追加した負の閉路があるグラフの入力例では有向グラフを使っています。

そのため、この記事では入力例は有向グラフに統一してあります。


V:頂点数\\E:辺数\\頂点f_iから頂点t_iに向かう重みc_iの有向辺が存在する(1\leq{i}\leq{E})


V\,E\\f_1\,t_1\,c_1\\:\\f_E\,t_E\,c_E

入力例1

f:id:kuruton456:20201031230032p:plain


7 20
0 1 2
1 0 2
0 2 5
2 0 5
1 2 4
2 1 4
1 3 6
3 1 6
1 4 10
4 1 10
2 3 2
3 2 2
3 5 1
5 3 1
4 5 3
5 4 3
4 6 5
6 4 5
5 6 9
6 5 9


入力例2

f:id:kuruton456:20201029231744p:plain


6 7
0 1 2
0 3 4
1 2 3
2 3 5
2 5 2
4 2 -4
5 4 1


負の閉路が無い場合のコード

INF = float('inf')

V, E = map(int,input().split())    # Vは頂点数、Eは辺数
es = []    # 辺
for _ in range(E):
    # 頂点fromから頂点toへのコストcostの辺
    f, t, c = map(int,input().split())
    es.append([f, t, c])

d = [INF] * V    # 最短距離

# s番目の頂点から各頂点への最短経路を求める
def shortest_path(s):
    d[s] = 0
    while True:
        update = False
        for i in range(E):
            e = es[i]
            if d[e[0]] != INF and d[e[1]] > d[e[0]] + e[2]:
                d[e[1]] = d[e[0]] + e[2]
                update = True
        if not update:
            break

shortest_path(0)    # 頂点0を始点とする

print(d)

実行結果

入力例1の場合
[0, 2, 5, 7, 11, 8, 16]
入力例2の場合

負の閉路があるため無限ループに入ります。


負閉路を検出するコード

INF = float('inf')

V, E = map(int,input().split())    # Vは頂点数、Eは辺数
es = []    # 辺
for _ in range(E):
    # 頂点fromから頂点toへのコストcostの辺
    f, t, c = map(int,input().split())
    es.append([f, t, c])

# trueなら負の閉路が存在する
def find_negative_loop():
    d = [0] * V
    
    for i in range(V):
        for j in range(E):
            e = es[j]
            if d[e[1]] > d[e[0]] + e[2]:
                d[e[1]] = d[e[0]] + e[2]
                
                # n回目にも更新があるなら負の経路が存在する
                if i == V-1:
                    return True
    return False

find_negative_loop()

実行結果

入力例1の場合
True
入力例2の場合
False

最短距離を求め負閉路があれば検出するコード

蟻本には載っていませんでしたが、上記の二つを合わせたものを載せておきます。

このコードでは、負閉路を利用することで最短距離が無限に小さくなるような頂点との距離を-INFと出力します。

このコードも計算量はO(|V|\times|E|)になります。

INF = float('inf')

V, E = map(int,input().split())    # Vは頂点数、Eは辺数
es = []    # 辺
for _ in range(E):
    # 頂点fromから頂点toへのコストcostの辺
    f, t, c = map(int,input().split())
    es.append([f, t, c])

def bellmanford(s):
    d = [INF] * V
    d[s] = 0
    
    for i in range(V):
        for j in range(E):
            e = es[j]
            if d[e[1]] > d[e[0]] + e[2]:
                d[e[1]] = d[e[0]] + e[2]
                
                # n回目にも更新があるなら負の経路が存在する
                if i == V-1:
                    d[e[1]] = -INF
                    while True:
                        update = False
                        for i in range(len(es)):
                            e = es[i]
                            if d[e[1]] != -INF and d[e[0]] == -INF:
                                d[e[1]] = -INF
                                update = True
                        if not update:
                            break
    return d

D = bellmanford(0)    # 頂点0を始点とする

print(D)

実行結果

入力例1の場合
[0, 2, 5, 7, 11, 8, 16]
入力例2の場合
[0, 2, -inf, -inf, -inf, -inf]