85 lines
2.0 KiB
Python
85 lines
2.0 KiB
Python
import random
|
|
from dinic_maxflow import max_flow as dinic_flow
|
|
|
|
|
|
def ek_flow(n, edges, s, t):
|
|
# Edmonds-Karp (BFS augmenting paths) reference
|
|
from collections import deque
|
|
cap = [[0]*n for _ in range(n)]
|
|
for u, v, c in edges:
|
|
cap[u][v] += c
|
|
flow = 0
|
|
while True:
|
|
parent = [-1]*n
|
|
parent[s] = s
|
|
q = deque([s])
|
|
while q and parent[t] == -1:
|
|
u = q.popleft()
|
|
for v in range(n):
|
|
if parent[v] == -1 and cap[u][v] > 0:
|
|
parent[v] = u
|
|
q.append(v)
|
|
if parent[t] == -1:
|
|
break
|
|
# bottleneck
|
|
aug = 10**18
|
|
v = t
|
|
while v != s:
|
|
u = parent[v]
|
|
aug = min(aug, cap[u][v])
|
|
v = u
|
|
v = t
|
|
while v != s:
|
|
u = parent[v]
|
|
cap[u][v] -= aug
|
|
cap[v][u] += aug
|
|
v = u
|
|
flow += aug
|
|
return flow
|
|
|
|
|
|
def test_simple():
|
|
assert dinic_flow(4, [(0,1,10),(1,3,10),(0,2,10),(2,3,10)], 0, 3) == 20
|
|
|
|
|
|
def test_single():
|
|
assert dinic_flow(2, [(0,1,5)], 0, 1) == 5
|
|
|
|
|
|
def test_zero_cap():
|
|
assert dinic_flow(2, [(0,1,0)], 0, 1) == 0
|
|
|
|
|
|
def test_disconnected():
|
|
assert dinic_flow(4, [(0,1,5)], 0, 3) == 0
|
|
|
|
|
|
def test_bottleneck():
|
|
assert dinic_flow(3, [(0,1,10),(1,2,3)], 0, 2) == 3
|
|
|
|
|
|
def test_multi_edge():
|
|
assert dinic_flow(2, [(0,1,3),(0,1,4)], 0, 1) == 7
|
|
|
|
|
|
def test_random_vs_ek():
|
|
rng = random.Random(42)
|
|
for _ in range(300):
|
|
n = rng.randint(2, 7)
|
|
s, t = 0, n-1
|
|
m = rng.randint(0, 12)
|
|
edges = []
|
|
for _ in range(m):
|
|
u = rng.randrange(n)
|
|
v = rng.randrange(n)
|
|
c = rng.randint(0, 6)
|
|
edges.append((u, v, c))
|
|
a = dinic_flow(n, edges, s, t)
|
|
b = ek_flow(n, edges, s, t)
|
|
assert a == b, (edges, a, b)
|
|
|
|
|
|
def test_self_loop_ignored():
|
|
# self-loops shouldn't inflate flow
|
|
assert dinic_flow(2, [(0,0,100),(0,1,5)], 0, 1) == 5
|