91 lines
2.6 KiB
Python
91 lines
2.6 KiB
Python
import random
|
|
from mincost_maxflow import MinCostMaxFlow
|
|
|
|
|
|
def brute_min_cost(n, edges, s, t, maxf):
|
|
INF = float("inf")
|
|
g = [[] for _ in range(n)]
|
|
for (u, v, cap, cost) in edges:
|
|
fi = len(g[u]); ri = len(g[v])
|
|
g[u].append([v, cap, cost, ri])
|
|
g[v].append([u, 0, -cost, fi])
|
|
flow = 0; cost = 0
|
|
while flow < maxf:
|
|
dist = [INF]*n; prevv=[-1]*n; preve=[-1]*n
|
|
dist[s]=0
|
|
for _ in range(n):
|
|
upd=False
|
|
for u in range(n):
|
|
if dist[u]==INF: continue
|
|
for i,e in enumerate(g[u]):
|
|
v,cap,cc,_=e
|
|
if cap>0 and dist[u]+cc<dist[v]:
|
|
dist[v]=dist[u]+cc; prevv[v]=u; preve[v]=i; upd=True
|
|
if not upd: break
|
|
if dist[t]==INF: break
|
|
addf=maxf-flow
|
|
v=t
|
|
while v!=s:
|
|
u=prevv[v]; e=g[u][preve[v]]; addf=min(addf,e[1]); v=u
|
|
v=t
|
|
while v!=s:
|
|
u=prevv[v]; ei=preve[v]; e=g[u][ei]
|
|
e[1]-=addf; g[v][e[3]][1]+=addf; cost+=addf*e[2]; v=u
|
|
flow+=addf
|
|
return flow, cost
|
|
|
|
|
|
def test_single_path():
|
|
m = MinCostMaxFlow(4)
|
|
m.add_edge(0,1,2,3); m.add_edge(1,2,2,4); m.add_edge(2,3,2,1)
|
|
f,c = m.min_cost_flow(0,3,2)
|
|
assert (f,c)==(2, 2*(3+4+1)), (f,c)
|
|
|
|
|
|
def test_two_parallel_paths():
|
|
m = MinCostMaxFlow(4)
|
|
m.add_edge(0,1,1,1); m.add_edge(1,3,1,1)
|
|
m.add_edge(0,2,1,5); m.add_edge(2,3,1,5)
|
|
f,c = m.min_cost_flow(0,3,2)
|
|
assert (f,c)==(2, (1+1)+(5+5)), (f,c)
|
|
|
|
|
|
def test_negative_cost():
|
|
m = MinCostMaxFlow(3)
|
|
m.add_edge(0,1,1,5); m.add_edge(1,2,1,-3); m.add_edge(0,2,1,10)
|
|
f,c = m.min_cost_flow(0,2,1)
|
|
assert (f,c)==(1,2), (f,c)
|
|
|
|
|
|
def test_random_vs_brute():
|
|
random.seed(42)
|
|
for trial in range(400):
|
|
n = random.randint(2,5)
|
|
s=0; t=n-1
|
|
edges=[]
|
|
for u in range(n):
|
|
for v in range(n):
|
|
if u==v: continue
|
|
if random.random()<0.4:
|
|
edges.append((u,v,random.randint(0,3),random.randint(-3,5)))
|
|
maxf=random.randint(1,6)
|
|
m=MinCostMaxFlow(n)
|
|
for (u,v,cap,cost) in edges:
|
|
m.add_edge(u,v,cap,cost)
|
|
f1,c1=m.min_cost_flow(s,t,maxf)
|
|
f2,c2=brute_min_cost(n,edges,s,t,maxf)
|
|
assert (f1,c1)==(f2,c2), (trial,(f1,c1),(f2,c2))
|
|
|
|
|
|
def test_maxflow_cap():
|
|
m=MinCostMaxFlow(2)
|
|
m.add_edge(0,1,5,2)
|
|
f,c=m.min_cost_flow(0,1,10)
|
|
assert (f,c)==(5,10)
|
|
|
|
|
|
if __name__=="__main__":
|
|
test_single_path(); test_two_parallel_paths(); test_negative_cost()
|
|
test_random_vs_brute(); test_maxflow_cap()
|
|
print("ALL OK")
|