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

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

Pythonで理解する蟻本「2-5 Layout（POJ No.3169）」(p.104)

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

この記事は「プログラミングコンテストチャレンジブック第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)

Pythonで理解する蟻本「2-5 Conscription（POJ No.3723）」(p.103)

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

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

入力

N\,M\,R\\x_1\,y_1\,d_1\\:\\x_R\,y_R\,d_R

入力例



5 5 8
4 3 6831
1 3 4583
0 0 6592
0 1 3063
3 3 4975
1 3 2049
4 2 2104
2 2 781


解答

# クラスカル法のコード

class UnionFindTree():
    # n要素で初期化
    def __init__(self, n):
        self.n = n
        self.par = list(range(n))
        self.rank = [0] * n

    # 木の根を求める
    def find(self, x):
        if self.par[x] == x:
            return x
        else:
            self.par[x] = self.find(self.par[x])
            return self.par[x]

    # xとyの属する集合を併合
    def unite(self, x, y):
        x = self.find(x)
        y = self.find(y)
        if x == y:
            return

        if self.rank[x] < self.rank[y]:
            self.par[x] = y
        else:
            self.par[y] = x
            if self.rank[x] == self.rank[y]:
                self.rank[x] += 1

    # xとyが同じ集合に属するか否か
    def same(self, x, y):
        return self.find(x) == self.find(y)

def kruskal():
    es.sort()    # 辺のコストが小さい順にソートする
    UFT = UnionFindTree(V)    # Union-Findの初期化
    res = 0
    for i in range(E):
        e = es[i]
        if not UFT.same(e[1], e[2]):
            UFT.unite(e[1], e[2])
            res += e[0]
    return res


# ここから蟻本のコード

MAX_R = 50000

# 入力
N, M, R = map(int,input().split())
x = [0] * MAX_R
y = [0] * MAX_R
d = [0] * MAX_R

for i in range(R):
    x[i], y[i], d[i] = map(int,input().split())

V = N + M
E = R
es = []
for i in range(R):
    es.append([-d[i], N + x[i], y[i]])
print(10000 * (N + M) + kruskal())

Pythonで理解する蟻本「2-5 Roadblocks（POJ No.3255）」(p.102)

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

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

入力


N:交差点の数\\R:道の本数\\交差点s_iと交差点t_iの間に長さc_iの道が存在する(1\leq{i}\leq{R})


N\,R\\s_1\,t_1\,c_1\\:\\s_R\,t_R\,c_R

入力例



4 4
1 2 100
2 3 250
2 4 200
3 4 100


解答

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

MAX_N = 5000
MAX_R = 100000
INF = float('inf')

# 入力
N, R = map(int,input().split())
G = [[] for _ in range(MAX_N)]    # グラフの隣接リスト表現
for i in range(R):
    s, t, c = map(int,input().split())
    G[s-1].append([t-1, c])
    G[t-1].append([s-1, c])

que = []
dist = [INF] * MAX_N     # 最短経路
dist2 = [INF] * MAX_N    # 二番目の最短経路
dist[0] = 0
heapq.heappush(que, [0, 0])

while len(que) != 0:
    p = heapq.heappop(que)
    d = p[0]
    v = p[1]
    if dist2[v] < d:
        continue
    for i in range(len(G[v])):
        e = G[v][i]
        d2 = d + e[1]
        if dist[e[0]] > d2:
            dist[e[0]], d2 = d2, dist[e[0]]
            heapq.heappush(que, [dist[e[0]], e[0]])
        if dist2[e[0]] > d2 and dist[e[0]] < d2:
            dist2[e[0]] = d2
            heapq.heappush(que, [dist2[e[0]], e[0]])
print(dist2[N - 1])

Pythonで理解する蟻本「2-5 最小全域木問題２（クラスカル法）」(p.101)

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

この記事は「プログラミングコンテストチャレンジブック第2版」（蟻本）の
「2-5 最小全域木問題２（クラスカル法）」(p.101)
のコードを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


コード

# Union-Find木のコード

class UnionFindTree():
    # n要素で初期化
    def __init__(self, n):
        self.n = n
        self.par = list(range(n))
        self.rank = [0] * n

    # 木の根を求める
    def find(self, x):
        if self.par[x] == x:
            return x
        else:
            self.par[x] = self.find(self.par[x])
            return self.par[x]

    # xとyの属する集合を併合
    def unite(self, x, y):
        x = self.find(x)
        y = self.find(y)
        if x == y:
            return

        if self.rank[x] < self.rank[y]:
            self.par[x] = y
        else:
            self.par[y] = x
            if self.rank[x] == self.rank[y]:
                self.rank[x] += 1

    # xとyが同じ集合に属するか否か
    def same(self, x, y):
        return self.find(x) == self.find(y)


# ここから蟻本のコード

V, E = map(int,input().split())    # Vは頂点数、Eは辺数
es = []

for _ in range(E):
    s, t, c = map(int,input().split())
    es.append([c, s, t])    # c(辺のコスト)でソートできるように最初の要素にしておく

def kruskal():
    es.sort()    # 辺のコストが小さい順にソートする
    UFT = UnionFindTree(V)    # Union-Findの初期化
    res = 0
    for i in range(E):
        e = es[i]
        if not UFT.same(e[1], e[2]):
            UFT.unite(e[1], e[2])
            res += e[0]
    return res

res = kruskal()

print(res)

実行結果

17

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が存在します。

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