Python
AtCoder
Pypy

競技プログラミングで使う有名グラフアルゴリズムまとめ


0. はじめに

AtCoderなどでは、グラフを扱った問題が多く出るが、その度に一から実装していると時間が掛かりすぎてしまうため、有名なものをあらかじめ持っておく必要がありそう。そこで、Pythonを用いて、ダイクストラ法、ベルマンフォード法、プリム法、クラスカル法、ワーシャルフロイド法を実装した。

コメント、意見等ある方は是非! お待ちしてます!


1. ダイクストラ法


1.1. ダイクストラ法(defaultdictで実装)

defaultdictで実装すると、リストで実装するよりも、ノード数$N$が大きい際には高速に動作する。ただし、経路復元の関数は、うまく書けなかった......。

import collections

import heapq

class Dijkstra():
def __init__(self):
self.e = collections.defaultdict(list)

def add(self, u, v, d, directed=False):
"""
0-indexedでなくてもよいことに注意
u = from, v = to, d = cost
directed = Trueなら、有向グラフである
"""

if directed is False:
self.e[u].append([v, d])
self.e[v].append([u, d])
else:
self.e[u].append([v, d])

def delete(self, u, v):
self.e[u] = [_ for _ in self.e[u] if _[0] != v]
self.e[v] = [_ for _ in self.e[v] if _[0] != u]

def Dijkstra_search(self, s):
"""
0-indexedでなくてもよいことに注意
s = 始点
return: 始点から各点までの最短経路
"""

d = collections.defaultdict(lambda: float('inf'))
d[s] = 0
q = []
heapq.heappush(q, (0, s))
v = collections.defaultdict(bool)
while len(q):
k, u = heapq.heappop(q)
if v[u]: # DefaulfではFalseとなる
continue

for uv, ud in self.e[u]:
if v[uv]:
continue
vd = k + ud
if d[uv] > vd:
d[uv] = vd
heapq.heappush(q, (vd, uv))

return d


ダイクストラ法(defaultdictで実装)の使用例

SoundHound2018 D - Saving Snuukに対する自分の解答

# クラス部分は上と全く一緒

N, M, S, T = map(int, input().split())
UVAB = [list(map(int, input().split())) for i in range(M)]
graph1 = Dijkstra()
graph2 = Dijkstra()

for u, v, a, b in UVAB:
graph1.add(u, v, a)
graph2.add(u, v, b)

result_a = graph1.search(S)
result_b = graph2.search(T)
ans_list = [float('inf')]
for i in range(N, 0, -1):
cur = result_a[i] + result_b[i]
if cur > ans_list[-1]:
cur = ans_list[-1]
ans_list.append(cur)

for i in range(len(ans_list[1:])):
print(10**15-ans_list[-(i+1)])


1.2. ダイクストラ法(リストで実装)

上にも書いた通り、動作が遅く、TLEすることも十分考えられる。しかし、経路復元が必要なときはこっちを使わざるを得ない。

import collections

import heapq

class Dijkstra():
def __init__(self, N):
self.N = N
self.e = [[float('inf') for i in range(self.N)] for i in range(N)]

def add(self, u, v, d, directed=False):
"""
0-indexedでなくてもよいことに注意
u = from, v = to, d = cost
directed = Trueなら、有向グラフである
"""

if directed is False:
self.e[u][v] = d
self.e[v][u] = d
else:
self.e[u][v] = d

def delete(self, u, v):
self.e[u] = [_ for _ in self.e[u] if _[0] != v]
self.e[v] = [_ for _ in self.e[v] if _[0] != u]

def Dijkstra_search(self, s):
"""
0-indexedであることに注意
s = 始点
return: 始点から各点までの最短経路
"""

d = [float('inf') for i in range(self.N)]
d[s] = 0
q = []
heapq.heappush(q, (0, s))
v = collections.defaultdict(bool)
while len(q):
k, u = heapq.heappop(q)
if v[u]:
continue

for uv, ud in self.e[u]:
if v[uv]:
continue
vd = k + ud
if d[uv] > vd:
d[uv] = vd
heapq.heappush(q, (vd, uv))

return d

def getDijkstraPath(self, s, t):
# sからtへの最短経路の経路復元
prev = [s] * self.N # 最短経路の直前の頂点
d = [float("inf")] * self.N
used = [False] * self.N
d[s] = 0

while True:
v = -1
for i in range(self.N):
if (not used[i]) and (v == -1):
v = i
elif (not used[i]) and d[i] < d[v]:
v = i
if v == -1:
break
used[v] = True

for i in range(self.N):
if d[i] > d[v] + self.e[v][i]:
d[i] = d[v] + self.e[v][i]
prev[i] = v

path = [t]
while prev[t] != s:
path.append(prev[t])
prev[t] = prev[prev[t]]
path.append(s)
path = path[::-1]
return path


2. ベルマンフォード法

ダイクストラ法と同様に、defaultdictで実装する。リストでの実装ももちろん可能。

import collections

class BellmanFord():
def __init__(self):
self.edges = []

def add(self, u, v, d, directed=False):
"""
u = from, v = to, d = cost
directed = Trueのとき、有向グラフである。
"""

if directed is False:
self.edges.append([u, v, d])
self.edges.append([v, u, d])
else:
self.edges.append([u, v, d])

def BellmanFord_search(self, s):
"""
次のBellmanFord_negative_boolまで考慮した時には、0-indexedであることに注意
s = 始点
return: d[i] 始点sから各点iまでの最短経路
"""

d = collections.defaultdict(lambda: float('inf'))
d[s] = 0
numEdges = len(self.edges)
while True:
update = False
for i in range(numEdges):
e = self.edges[i]
# e: 辺iについて [from,to,cost]
if d[e[0]] != float("inf") and d[e[1]] > d[e[0]] + e[2]:
d[e[1]] = d[e[0]] + e[2]
update = True
if not update:
break
return d

def BellmanFord_negative_bool(self, numNodes, zero_indexed=False):
# 負の閉路の検出, Trueなら負の閉路が存在する
d = collections.defaultdict(lambda: float('inf'))
d[0] = 0
numEdges = len(self.edges)
for i in range(numNodes):
for j in range(numEdges):
e = self.edges[j]
if d[e[1]] > d[e[0]] + e[2]:
d[e[1]] = d[e[0]] + e[2]
if i == numNodes-1:
return True
return False


3. プリム法

import heapq

class Prim():
# 無向グラフであるという前提に注意
def __init__(self, N):
self.edge = [[] for i in range(N)]
self.N = N

def add(self, u, v, d):
"""
u = from, v = to, d = cost
0-indexedであることに注意、graph.add(u-1, v-1)とする必要がある
"""

self.edge[u].append([d, v]) # コスト、e_toとなっていることに注意
self.edge[v].append([d, u])

def delete(self, u, v):
self.edge[u] = [_ for _ in self.edge[u] if _[0] != v]
self.edge[v] = [_ for _ in self.edge[v] if _[0] != u]

def Prim(self):
"""
return: 最小全域木のコストの和
"""

used = [True] * self.N # True:不使用
edgelist = []
for e in self.edge[0]:
heapq.heappush(edgelist, e)
used[0] = False
res = 0
while len(edgelist) != 0:
minedge = heapq.heappop(edgelist)
if not used[minedge[1]]:
continue
v = minedge[1]
used[v] = False
for e in self.edge[v]:
if used[e[1]]:
heapq.heappush(edgelist, e)
res += minedge[0]
return res


4. クラスカル法

class Kruskal_UnionFind():

# 無向グラフであるという前提に注意
def __init__(self, N):
self.edges = []
self.rank = [0] * N
self.par = [i for i in range(N)]
self.counter = [1] * N

def add(self, u, v, d):
"""
u = from, v = to, d = cost
"""

self.edges.append([u, v, d])

def find(self, x):
if self.par[x] == x:
return x
else:
self.par[x] = self.find(self.par[x])
return self.par[x]

def unite(self, x, y):
x = self.find(x)
y = self.find(y)
if x != y:
z = self.counter[x] + self.counter[y]
self.counter[x], self.counter[y] = z, z
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

def size(self, x):
x = self.find(x)
return self.counter[x]

def same(self, x, y):
return self.find(x) == self.find(y)

def Kruskal(self):
"""
return: 最小全域木のコストの和
"""

edges = sorted(self.edges, key=lambda x: x[2]) # costでself.edgesをソートする
res = 0
for e in edges:
if not self.same(e[0], e[1]):
self.unite(e[0], e[1])
res += e[2]
return res


クラスカル法の使用例

AtCoder Beginner Contest 065 D - Built?に対する自分の解答

# クラス部分は上と一緒

N = int(input())
XY = [[i] + list(map(int, input().split())) for i in range(N)]

graph = Kruskal_UnionFind(N)
XY = sorted(XY, key=lambda x: x[1])
X_costs = [[XY[i-1][0], XY[i][0], abs(XY[i-1][1] - XY[i][1])] for i in range(1, N)]
XY = sorted(XY, key=lambda x: x[2])
Y_costs = [[XY[i-1][0], XY[i][0], abs(XY[i-1][2] - XY[i][2])] for i in range(1, N)]

for i in range(N-1):
x0, x1, d = X_costs[i]
graph.add(x0, x1, d)
y0, y1, d = Y_costs[i]
graph.add(y0, y1, d)

print(graph.Kruskal())


5. ワーシャルフロイド法

class WarshallFloyd():

def __init__(self, N):
self.N = N
self.d = [[float("inf") for i in range(N)]
for i in range(N)] # d[u][v] : 辺uvのコスト(存在しないときは'inf')

def add(self, u, v, c, directed=False):
"""
0-indexedであることに注意
u = from, v = to, c = cost
directed = Trueなら、有向グラフである
"""

if directed is False:
self.d[u][v] = c
self.d[v][u] = c
else:
self.d[u][v] = c

def WarshallFloyd_search(self):
# これを d[i][j]: iからjへの最短距離 にする
# 本来無向グラフでのみ全域木を考えるが、二重辺なら有向でも行ける
for k in range(self.N):
for i in range(self.N):
for j in range(self.N):
self.d[i][j] = min(
self.d[i][j], self.d[i][k] + self.d[k][j])
for i in range(self.N):
self.d[i][i] = 0
return self.d


ワーシャルフロイド法の使用例

AtCoder Beginner Contest 012 D - バスと避けられない運命に対する自分の解答

# クラス部分は上と一緒

N, M = map(int, input().split())
ABT = [list(map(int, input().split())) for i in range(M)]
graph = WarshallFloyd(N)
for a, b, t in ABT:
graph.add(a-1, b-1, t)

d = graph.WarshallFloyd_search()
ans = sorted([[i, max(d[i])] for i in range(N)], key=lambda x: x[1])
print(ans[0][1])