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

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

Pythonで理解する蟻本「2-5 最小全域木問題１（プリム法）」(p.100)

Pythonで理解する蟻本（プログラミングコンテストチャレンジブック）

この記事は「プログラミングコンテストチャレンジブック第2版」（蟻本）の
「2-5 最小全域木問題１（プリム法）」(p.100)
のコードをPythonで書き直したものとなっています。

入力


V:頂点数\\E:辺数\\頂点s_iと頂点t_iの間に重みc_iの無向辺が存在する(1\leq{i}\leq{E})


V\,E\\s_1\,t_1\,c_1\\:\\s_E\,t_E\,c_E

入力例

f:id:kuruton456:20201201145758p:plain


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


蟻本に載っているコード（O(|V|^2)）

INF = float('inf')

V, E = map(int,input().split())    # Vは頂点数、Eは辺数
cost = [[INF] * V for _ in range(V)]    # cost[u][v]は辺e=(u,v)のコスト(存在しない場所はINF)
mincost = [INF] * V    # 集合Xからの辺の最終コスト
used = [False] * V    # 頂点iがXに含まれているか

for _ in range(E):
    s, t, c = map(int,input().split())
    cost[s][t] = c
    cost[t][s] = c

def prim():
    mincost[0] = 0
    res = 0
    
    while True:
        v = -1
        # Xに属さない頂点のうちXからの辺のコストが最小になる頂点を探す
        for u in range(V):
            if not used[u] and (v == -1 or mincost[u] < mincost[v]):
                v = u
        
        if v == -1:
            break
        used[v] = True    # 頂点vをXに追加
        res += mincost[v]    # 辺のコストを加える
        
        for u in range(V):
            mincost[u] = min(mincost[u], cost[v][u])
        
    return res

res = prim()

print(res)

実行結果

17

優先度付きキューを使った解法（O(|E| log |V|)）

# 優先度付きキューをインポートする
import heapq

INF = float('inf')

V, E = map(int,input().split())    # Vは頂点数、Eは辺数
cost = [[INF] * V for _ in range(V)]    # cost[u][v]は辺e=(u,v)のコスト(存在しない場所はINF)
used = [False] * V    # 頂点iがXに含まれているか

for _ in range(E):
    s, t, c = map(int,input().split())
    cost[s][t] = c
    cost[t][s] = c

def prim():
    res = 0
    q = [[0, 0]]
    
    while True:
        v = -1
        MIN = INF
        # Xに属さない頂点のうちXからの辺のコストが最小になる頂点を探す
        while len(q) != 0:
            c, u = heapq.heappop(q)
            if not used[u] and (v == -1 or c < MIN):
                v = u
                MIN = c
                break
        
        if v == -1:
            break
        used[v] = True    # 頂点vをXに追加
        res += c    # 辺のコストを加える
        
        for u in range(V):
            if cost[v][u] != INF:
                heapq.heappush(q, [cost[v][u], u])
        
    return res

res = prim()

print(res)

実行結果

17

Pythonで理解する蟻本「2-5 経路復元」(p.98)

Pythonで理解する蟻本（プログラミングコンテストチャレンジブック）

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

入力

ダイクストラ法は負の辺があると使えないため、負閉路の存在を考える必要はありません。

そのため、この記事では入力例は蟻本のテストケースのみにし、無向辺グラフとして入力を受け取っています。


V:頂点数\\E:辺数\\頂点s_iと頂点t_iの間に重みc_iの無向辺が存在する(1\leq{i}\leq{E})


V\,E\\s_1\,t_1\,c_1\\:\\s_E\,t_E\,c_E

入力例

f:id:kuruton456:20201029232302p:plain


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


隣接リストを用いたコード（O(|V|^2)）

INF = float('inf')

prev = []

V, E = map(int,input().split())    # Vは頂点数、Eは辺数
cost = [[INF] * V for _ in range(V)]    # cost[u][v]は辺e=(u,v)のコスト(存在しない場所はINF)
for _ in range(E):
    s, t, c = map(int,input().split())
    cost[s][t] = c
    cost[t][s] = c


d = [INF for _ in range(V)]    # 最短距離
used = [False] * V    # すでに使われたかのフラグ
prev = [-1] * V    # 最短路の直前の頂点

# 始点sから各頂点への最短距離を求める
def dijkstra(s):
    d[s] = 0
    
    while True:
        v = -1
        for u in range(V):
            if not used[u] and (v == -1 or d[u] < d[v]):
                v = u
        
        if v == -1:
            break
        used[v] = True
        
        for u in range(V):
            if d[u] > d[v] + cost[v][u]:
                d[u] = d[v] + cost[v][u]
                prev[u] = v

# 頂点tへの最短路
def get_path(t):
    path = []
    while t != -1:
        path.append(t)    # tがsになるまでprev[t]を辿っていく
        t = prev[t]
    # このままだとt->sの順になっているので逆順にする
    path.reverse()
    return path

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

path = get_path(6)    # 頂点0から頂点6への最短経路

print(path)

実行結果

[0, 2, 3, 5, 4, 6]

頂点0→2→3→5→4→6が最短経路となる。
f:id:kuruton456:20201105235135p:plain

Pythonで理解する蟻本「2-5 全点対最短路問題（ワーシャル-フロイド法）」(p.98)

Pythonで理解する蟻本（プログラミングコンテストチャレンジブック）

この記事は「プログラミングコンテストチャレンジブック第2版」（蟻本）の
「2-5 全点対最短路問題（ワーシャル-フロイド法）」(p.98)
のコードを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は辺数
d = [[INF] * V for _ in range(V)]    # d[u][v]は辺e=(u,v)のコスト(存在しない場合はINF、ただしd[i][i]=0とする)
for i in range(V):
    d[i][i] = 0
for _ in range(E):
    f, t, c = map(int,input().split())
    d[f][t] = c

def warshall_floyd():
    for k in range(V):
        for i in range(V):
            for j in range(V):
                d[i][j] = min(d[i][j], d[i][k] + d[k][j])

warshall_floyd()

print(d)

実行結果

上記のコードは実際に実行すると

[[0, 1, 2], [3, 4, 5], [6, 7, 8]]

のように一行で表示されます。

しかし、この形式はとても見にくいため下の実行例では

[[0, 1, 2],
 [3, 4, 5],
 [6, 7, 8]]

というように改行しています。

実際の実行結果とは異なりますが、リストの中身自体は同じなので問題ありません。

入力例1の場合
[[0,  2,  5,  7,  11, 8,  16],
 [2,  0,  4,  6,  10, 7,  15],
 [5,  4,  0,  2,  6,  3,  11],
 [7,  6,  2,  0,  4,  1,  9],
 [11, 10, 6,  4,  0,  3,  5],
 [8,  7,  3,  1,  3,  0,  8],
 [16, 15, 11, 9,  5,  8,  0]]
入力例2の場合
[[0,   2,   4,   4, 8,   6],
 [inf, 0,   2,   7, 6,   4],
 [inf, inf, -1,  4, 3,   1],
 [inf, inf, inf, 0, inf, inf],
 [inf, inf, -5,  0, -1,  -3],
 [inf, inf, -4,  1, 0,   -2]]


頂点iから頂点jに到達できない場合はd[i][j] = infとなり、負の閉路が存在する場合はd[i][i]が負となる頂点iが存在します。

これは上の実行例でも確かに成り立っています。

Pythonで理解する蟻本「2-5 単一始点最短路問題２（ダイクストラ法）」(p.96)

Pythonで理解する蟻本（プログラミングコンテストチャレンジブック）

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

入力

ダイクストラ法は負の辺があると使えないため、負閉路の存在を考える必要はありません。

そのため、この記事では入力例は蟻本のテストケースのみにし、無向辺グラフとして入力を受け取っています。


V:頂点数\\E:辺数\\頂点s_iと頂点t_iの間に重みc_iの無向辺が存在する(1\leq{i}\leq{E})


V\,E\\s_1\,t_1\,c_1\\:\\s_E\,t_E\,c_E

入力例

f:id:kuruton456:20201029232302p:plain


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


隣接リストを用いたコード（O(|V|^2)）

INF = float('inf')

V, E = map(int,input().split())    # Vは頂点数、Eは辺数
cost = [[INF] * V for _ in range(V)]    # cost[u][v]は辺e=(u,v)のコスト(存在しない場所はINF)
for _ in range(E):
    s, t, c = map(int,input().split())
    cost[s][t] = c
    cost[t][s] = c
d = [INF for _ in range(V)]    # 最短距離
used = [False] * V    # すでに使われたかのフラグ

# 始点sからの最短距離のリストを返す
def dijkstra(s):
    d[s] = 0
    
    while True:
        v = -1
        # まだ使われていない頂点のうち距離が最小のものを探す
        for u in range(V):
            if not used[u] and (v == -1 or d[u] < d[v]):
                v = u
        
        if v == -1:
            break
        used[v] = True
        
        for u in range(V):
            d[u] = min(d[u], d[v] + cost[v][u])

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

print(d)

実行結果

[0, 2, 5, 7, 11, 8, 16]

優先度付きキューを使った解法（O(|E| log |V|)）

# 優先度付きキューをインポートする
import heapq

INF = float('inf')

V, E = map(int,input().split())    # Vは頂点数、Eは辺数
G = [[] for _ in range(V)]
for _ in range(E):
    s, t, c = map(int,input().split())
    G[s].append([t, c])
    G[t].append([s, c])
d = [INF for _ in range(V)]

def dijkstra(s):
    d[s] = 0
    que = []
    heapq.heappush(que, (0, s))
    
    while len(que) != 0:
        p = heapq.heappop(que)
        v = p[1]
        if d[v] < p[0]:
            continue
        for i in range(len(G[v])):
            e = G[v][i]
            if d[e[0]] > d[v] + e[1]:
                d[e[0]] = d[v] + e[1]
                heapq.heappush(que, (d[e[0]], e[0]))

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

print(d)

実行結果

[0, 2, 5, 7, 11, 8, 16]

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

Pythonで理解する蟻本（プログラミングコンテストチャレンジブック）

この記事は「プログラミングコンテストチャレンジブック第2版」（蟻本）の
「2-5 単一始点最短路問題１（ベルマンフォード法）」(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]

Pythonで理解する蟻本「2-5 二部グラフ判定」(p.93)

Pythonで理解する蟻本（プログラミングコンテストチャレンジブック）

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

入力

蟻本の入力例ではグラフが書かれていただけだったので、この記事ではAtCoderでよくあるような入力形式で入力例を書き直しておきました。

そのため、グラフの入力を受け取る部分は蟻本に載っているコードとは異なります。

この入力の受け取り方は『Pythonで理解する蟻本「2-5 隣接リスト」(p.91)』に載っているのでそちらを参考にしてください。


n:頂点数\\m:辺数\\2点s_i,\,t_iの間に無向辺が存在する(1\leq{i}\leq{m})


n\\m\\s_1\,t_1\\:\\s_m\,t_m

入力例

入力例1



3
3
0 1
0 2
1 2

入力例1



4
4
0 1
0 3
1 2
2 3


解答

# DFSを使うため再帰限界に引っかからないようにしておく
import sys
sys.setrecursionlimit(4100000)

MAX_V = 1000

# 入力
V = int(input())    # 頂点数
E = int(input())    # 辺数
G = [[] for _ in range(MAX_V)]    # グラフ
for i in range(E):
    # sからtへの辺を張る
    s, t = map(int,input().split())
    G[s].append(t)
    G[t].append(s)

color = [0] * MAX_V    # 頂点iの色(1 or -1)

# 頂点を1と-1で塗っていく
def dfs(v, c):
    color[v] = c    # 頂点vをcで塗る
    for i in range(len(G[v])):
        # 隣接している頂点が同じ色ならFalse
        if color[G[v][i]] == c:
            return False
        # 隣接している頂点がまだ塗られていないなら-cで塗る
        if color[G[v][i]] == 0 and not dfs(G[v][i], -c):
            return False
    # すべての頂点を塗れたらTrue
    return True

for i in range(V):
    if color[i] == 0:
        # まだ頂点iが塗られていなければ1で塗る
        if not dfs(i, 1):
            print("No")
            sys.exit()
print("Yes")

Pythonで理解する蟻本「2-5 隣接リスト」(p.91)

Pythonで理解する蟻本（プログラミングコンテストチャレンジブック）

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

隣接リストのコード

例１

MAX_V = 10 ** 8
G = [0] * MAX_V

'''
辺に属性がある場合
edge = [to, cost]
G = [edge for _ in range(MAX_V)]
'''

V = int(input())
E = int(input())
for i in range(E):
    # sからtへの辺を張る
    s = int(input())
    t = int(input())
    G[s].append(t)
    # G[t].append(s)    無向グラフの場合はさらにtにsへの辺を張る

'''
グラフの操作
'''


例２のコードは翻訳できませんでした…

翻訳できる方がいらっしゃったら僕のツイッターにコードを送って頂けると助かりますm(_ _)m