121 lines
4.0 KiB
Python
121 lines
4.0 KiB
Python
"""Minimum-cost maximum-flow (successive shortest augmenting path with potentials).
|
|
|
|
Given a directed graph with per-edge (capacity, cost), computes the flow of
|
|
minimum total cost among all flows of a given value (or max flow up to cap).
|
|
|
|
Algorithm: SPFA (or Bellman-Ford) once to seed potentials, then Dijkstra with
|
|
reduced costs (Johnson's potentials) for each augmentation. Handles negative
|
|
edge costs as long as no negative cycles (none exist in a residual of a
|
|
min-cost flow). Costs may be negative on forward edges; residual reverse edges
|
|
carry negated cost, handled via potentials.
|
|
|
|
API:
|
|
mcmf = MinCostMaxFlow(n)
|
|
mcmf.add_edge(u, v, cap, cost)
|
|
flow, cost = mcmf.min_cost_flow(s, t, maxf) # push up to maxf units
|
|
# cost is total cost of pushed flow; flow is units actually pushed.
|
|
|
|
Invariants:
|
|
1. Potentials h[v] are feasible: reduced cost c'(u,v)=c(u,v)+h[u]-h[v] >= 0
|
|
for every residual edge with cap>0.
|
|
2. Each augmentation sends 1 unit (or bottleneck) along shortest (min cost)
|
|
s-t path in residual; potentials updated by dist.
|
|
3. Total cost accumulated equals sum over augmentations of path cost.
|
|
"""
|
|
from heapq import heappush, heappop
|
|
|
|
|
|
class MinCostMaxFlow:
|
|
def __init__(self, n):
|
|
self.n = n
|
|
self.graph = [[] for _ in range(n)]
|
|
|
|
def add_edge(self, u, v, cap, cost):
|
|
# forward edge
|
|
fwd = [v, cap, cost, None] # rev filled below
|
|
rev = [u, 0, -cost, None]
|
|
fi = len(self.graph[u])
|
|
ri = len(self.graph[v])
|
|
fwd[3] = ri
|
|
rev[3] = fi
|
|
self.graph[u].append(fwd)
|
|
self.graph[v].append(rev)
|
|
return fi
|
|
|
|
def _bellman(self, s):
|
|
INF = float("inf")
|
|
dist = [INF] * self.n
|
|
dist[s] = 0
|
|
# V-1 relaxations
|
|
for _ in range(self.n):
|
|
updated = False
|
|
for u in range(self.n):
|
|
if dist[u] == INF:
|
|
continue
|
|
for e in self.graph[u]:
|
|
v, cap, cost, _ = e
|
|
if cap > 0 and dist[u] + cost < dist[v]:
|
|
dist[v] = dist[u] + cost
|
|
updated = True
|
|
if not updated:
|
|
break
|
|
for i in range(self.n):
|
|
if dist[i] == INF:
|
|
dist[i] = 0
|
|
return dist
|
|
|
|
def min_cost_flow(self, s, t, maxf):
|
|
n = self.n
|
|
INF = float("inf")
|
|
h = self._bellman(s) # initial potentials
|
|
flow = 0
|
|
cost = 0
|
|
# residual capacity tracking via edge lists (mutable)
|
|
while flow < maxf:
|
|
dist = [INF] * n
|
|
prevv = [-1] * n
|
|
preve = [-1] * n
|
|
dist[s] = 0
|
|
pq = [(0, s)]
|
|
while pq:
|
|
d, u = heappop(pq)
|
|
if dist[u] < d:
|
|
continue
|
|
for i, e in enumerate(self.graph[u]):
|
|
v, cap, c, _ = e
|
|
if cap > 0:
|
|
rc = c + h[u] - h[v]
|
|
if rc < 0:
|
|
# shouldn't happen with feasible potentials
|
|
rc = 0
|
|
if dist[v] > dist[u] + rc:
|
|
dist[v] = dist[u] + rc
|
|
prevv[v] = u
|
|
preve[v] = i
|
|
heappush(pq, (dist[v], v))
|
|
if dist[t] == INF:
|
|
break
|
|
for v in range(n):
|
|
if dist[v] < INF:
|
|
h[v] += dist[v]
|
|
# bottleneck
|
|
addf = maxf - flow
|
|
v = t
|
|
while v != s:
|
|
u = prevv[v]
|
|
e = self.graph[u][preve[v]]
|
|
addf = min(addf, e[1])
|
|
v = u
|
|
v = t
|
|
while v != s:
|
|
u = prevv[v]
|
|
ei = preve[v]
|
|
e = self.graph[u][ei]
|
|
e[1] -= addf
|
|
revi = e[3]
|
|
self.graph[v][revi][1] += addf
|
|
cost += addf * e[2]
|
|
v = u
|
|
flow += addf
|
|
return flow, cost
|