"""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