Cost scaling min-cost flow algorithm in Python
Scale costs and refine epsilon-optimal flows via push-relabel style operations.
- File
cost_scaling.py, 125 lines- Level
- I0 not checked yet
- Toolchain
- Python 3.14, pytest 9
- Test framework
- pytest
- Run the tests
python3 -m pytest
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/cost-scaling/. The authors are listed in prov/ZWH6V8.jsonl.
from collections import deque
from typing import Callable, Optional
def cost_scaling(
n: int,
edges: list[tuple[int, int, int, int]],
s: int,
t: int,
on_refine: Optional[Callable[[int, int, int, int], None]] = None,
) -> tuple[int, int, list[int]]:
"""Minimum-cost maximum s-t flow: a maximum flow, then Goldberg-Tarjan cost scaling.
The network has the vertices 0 .. n-1 and the directed edges (u, v, cap, cost) of
the list, with integer cap >= 0 and integer cost of any sign. Parallel,
antiparallel and zero-capacity edges are allowed; each edge keeps its own flow.
A self-loop is full if its cost is negative and empty otherwise. Requires s != t.
Costs are multiplied by n + 1 inside, so that 1-optimality means optimality.
The optional callback gets (eps, pushes, relabels, cost) after each refine, with
eps in the scaled costs and cost the current cost in the original costs.
Returns (value, cost, flow): flow[i] is the flow on edges[i]; among all maximum
flows this one has the least total cost.
"""
if s == t:
raise ValueError("source and sink must differ")
head: list[int] = [] # arc 2i is edge i forward, arc 2i + 1 is edge i backward
res: list[int] = []
cost: list[int] = []
adj: list[list[int]] = [[] for _ in range(n)]
loops: list[int] = []
alpha = n + 1
for i, (u, v, cap, w) in enumerate(edges):
if not isinstance(cap, int) or not isinstance(w, int):
raise TypeError("capacities and costs must be integers")
if cap < 0:
raise ValueError("capacities must be nonnegative")
head.append(v)
res.append(cap)
cost.append(w * alpha)
head.append(u)
res.append(0)
cost.append(-w * alpha)
if u == v:
loops.append(i)
continue
adj[u].append(2 * i)
adj[v].append(2 * i + 1)
for i in loops:
if edges[i][3] < 0:
res[2 * i + 1], res[2 * i] = edges[i][2], 0
value = 0
while True: # @line L1
parent = [-1] * n
parent[s] = -2
queue = deque([s])
while queue and parent[t] == -1:
u = queue.popleft()
for a in adj[u]:
if res[a] > 0 and parent[head[a]] == -1:
parent[head[a]] = a
queue.append(head[a])
if parent[t] == -1:
break
delta = None
v = t
while v != s:
a = parent[v]
delta = res[a] if delta is None else min(delta, res[a])
v = head[a ^ 1]
v = t
while v != s:
a = parent[v]
res[a] -= delta
res[a ^ 1] += delta
v = head[a ^ 1]
value += delta
p = [0] * n
excess = [0] * n
eps = max([1] + [abs(c) for c in cost]) # @line L2
while True: # @line L3
eps = (eps + 1) // 2 # @line L4
pushes = relabels = 0
for v in range(n): # @line L5
for a in adj[v]:
w = head[a]
if res[a] > 0 and cost[a] + p[v] - p[w] < 0:
excess[v] -= res[a]
excess[w] += res[a]
res[a ^ 1] += res[a]
res[a] = 0
pushes += 1
active = deque(v for v in range(n) if excess[v] > 0)
it = [0] * n
while active: # @line L6
v = active.popleft()
while excess[v] > 0:
if it[v] < len(adj[v]):
a = adj[v][it[v]]
w = head[a]
if res[a] > 0 and cost[a] + p[v] - p[w] < 0: # @line L7
d = min(excess[v], res[a]) # @line L8
res[a] -= d
res[a ^ 1] += d
excess[v] -= d
excess[w] += d
pushes += 1
if excess[w] > 0 and excess[w] <= d:
active.append(w)
else:
it[v] += 1
else:
p[v] = max(p[head[a]] - cost[a] for a in adj[v] if res[a] > 0) - eps # @line L9
it[v] = 0
relabels += 1
if on_refine is not None:
cur = sum(edges[i][3] * res[2 * i + 1] for i in range(len(edges)))
on_refine(eps, pushes, relabels, cur)
if eps == 1: # @line L10
break
flow = [res[2 * i + 1] for i in range(len(edges))] # @line L11
total = sum(e[3] * f for e, f in zip(edges, flow))
return value, total, flow
test_cost_scaling.py 139 lines
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/cost-scaling/. The authors are listed in prov/ZWH6V8.jsonl.
import random
from itertools import product
import pytest
from cost_scaling import cost_scaling
S, A, B, C, D, T = range(6)
RUNNING = [
(S, A, 10, 2), (S, C, 10, 2), (A, B, 4, 6), (A, C, 2, 3), (A, D, 8, 1),
(C, D, 9, 2), (D, B, 6, 1), (B, T, 10, 1), (D, T, 10, 3),
]
def brute(n, edges, s, t):
"""Oracle: enumerate every integral flow; return (max value, least cost at that value)."""
best = None
for f in product(*[range(c + 1) for _, _, c, _ in edges]):
exc = [0] * n
for (u, v, _, _), x in zip(edges, f):
exc[u] -= x
exc[v] += x
if any(exc[v] for v in range(n) if v not in (s, t)):
continue
key = (-(-exc[s]), sum(x * e[3] for x, e in zip(f, edges)))
if best is None or key < best:
best = key
return -best[0], best[1]
def check(n, edges, s, t, value, cost, flow):
exc = [0] * n
for (u, v, c, _), f in zip(edges, flow):
assert 0 <= f <= c
exc[u] -= f
exc[v] += f
for v in range(n):
if v not in (s, t):
assert exc[v] == 0
assert -exc[s] == value
assert sum(f * e[3] for f, e in zip(flow, edges)) == cost
def test_running_example() -> None:
log = []
value, cost, flow = cost_scaling(6, RUNNING, S, T, lambda *a: log.append(a))
assert (value, cost) == (19, 126)
assert flow == [10, 9, 3, 0, 7, 9, 6, 9, 10]
assert [e for e, _, _, _ in log] == [21, 11, 6, 3, 2, 1]
assert log[0][3] == 126
check(6, RUNNING, S, T, value, cost, flow)
def test_no_path_zero_capacity_and_empty() -> None:
assert cost_scaling(3, [(0, 1, 5, 1), (2, 1, 5, 1)], 0, 2) == (0, 0, [0, 0])
assert cost_scaling(3, [(0, 1, 5, 1), (1, 2, 0, -9)], 0, 2) == (0, 0, [0, 0])
assert cost_scaling(2, [], 0, 1) == (0, 0, [])
def test_negative_cycle_away_from_the_path() -> None:
# a negative cycle not connected to s or t is still cancelled
edges = [(0, 1, 2, 1), (2, 3, 3, -2), (3, 2, 5, 1)]
value, cost, flow = cost_scaling(4, edges, 0, 1)
assert (value, cost, flow) == (2, 2 - 3, [2, 3, 3])
def test_negative_self_loop_is_filled() -> None:
value, cost, flow = cost_scaling(2, [(0, 1, 1, 0), (0, 0, 4, -1)], 0, 1)
assert (value, cost, flow) == (1, -4, [1, 4])
def test_antiparallel_edges_keep_their_own_flow() -> None:
edges = [(0, 1, 3, 1), (1, 2, 2, 4), (2, 1, 1, -5), (2, 3, 2, 1), (0, 2, 1, 1)]
value, cost, flow = cost_scaling(4, edges, 0, 3)
assert (value, cost) == brute(4, edges, 0, 3)
check(4, edges, 0, 3, value, cost, flow)
assert flow[1] >= 1 and flow[2] == 1 # the cycle a b a has cost -1 and is used
def test_cheaper_long_path() -> None:
edges = [(0, 3, 1, 10), (0, 1, 1, 1), (1, 2, 1, 1), (2, 3, 1, 1)]
value, cost, flow = cost_scaling(4, edges, 0, 3)
assert (value, cost) == (2, 13)
def test_bad_input() -> None:
with pytest.raises(ValueError):
cost_scaling(2, [(0, 1, 1, 0)], 0, 0)
with pytest.raises(ValueError):
cost_scaling(2, [(0, 1, -1, 0)], 0, 1)
with pytest.raises(TypeError):
cost_scaling(2, [(0, 1, 1.5, 0)], 0, 1)
def test_random_against_brute_force() -> None:
rng = random.Random(1989)
for _ in range(600):
n = rng.randint(2, 5)
m = rng.randint(0, 6)
edges = [(rng.randrange(n), rng.randrange(n), rng.randint(0, 3), rng.randint(-4, 6)) for _ in range(m)]
s, t = rng.sample(range(n), 2)
value, cost, flow = cost_scaling(n, edges, s, t)
assert (value, cost) == brute(n, edges, s, t)
check(n, edges, s, t, value, cost, flow)
def test_epsilon_sequence_and_large_costs() -> None:
log = []
edges = [(0, 1, 5, 1000), (0, 2, 5, -7), (2, 1, 5, 3), (1, 3, 9, 0)]
value, cost, flow = cost_scaling(4, edges, 0, 3, lambda *a: log.append(a))
assert (value, cost) == brute(4, edges, 0, 3)
eps = [e for e, _, _, _ in log]
assert eps[-1] == 1 and all(a > b for a, b in zip(eps, eps[1:]))
assert len(eps) <= (1000 * 5).bit_length() + 2
def test_final_flow_has_no_negative_cycle() -> None:
rng = random.Random(404)
for _ in range(200):
n = rng.randint(2, 7)
edges = [(rng.randrange(n), rng.randrange(n), rng.randint(0, 5), rng.randint(-9, 9)) for _ in range(rng.randint(0, 14))]
s, t = rng.sample(range(n), 2)
value, cost, flow = cost_scaling(n, edges, s, t)
dist = [0] * n
arcs = []
for (u, v, c, w), f in zip(edges, flow):
if f < c:
arcs.append((u, v, w))
if f > 0:
arcs.append((v, u, -w))
for _ in range(n):
for u, v, w in arcs:
if dist[u] + w < dist[v]:
dist[v] = dist[u] + w
assert all(dist[u] + w >= dist[v] for u, v, w in arcs)
check(n, edges, s, t, value, cost, flow)
What the next level needs
- the tests did not run in this build
Read the full page about Cost scaling min-cost flow algorithm