beam: min-cost max-flow (successive shortest path w/ potentials) + tests

wizard-work/beam/CHECK_mcmf

@@ -0,0 +1,4 @@
+#!/usr/bin/env bash
+set -euo pipefail
+cd "$(dirname "$(readlink -f "$0")")"
+python3 test_mincost_maxflow_small.py

wizard-work/beam/mincost_maxflow.py

@@ -0,0 +1,120 @@
+"""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

wizard-work/beam/test_mincost_maxflow.py

@@ -0,0 +1,90 @@
+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")

wizard-work/beam/test_mincost_maxflow_small.py

@@ -0,0 +1,10 @@
+import test_mincost_maxflow as t
+t.test_single_path()
+print("1")
+t.test_two_parallel_paths()
+print("2")
+t.test_negative_cost()
+print("3")
+t.test_maxflow_cap()
+print("4")
+print("SMALL OK")