Push-relabel algorithm in Python
Maintain a preflow and move excess locally using height labels until it becomes a flow.
- File
push_relabel.py, 76 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/push-relabel/. The authors are listed in prov/56SN6Q.jsonl.
from collections import deque
def push_relabel(n: int, edges: list[tuple[int, int, float]], s: int, t: int) -> tuple[float, list[float], set[int]]:
"""Maximum s-t flow by the FIFO push-relabel method of Goldberg and Tarjan.
The network is directed, with vertices 0 .. n-1 and edges (u, v, c), where
c >= 0 is an integer or a real number. Parallel and antiparallel edges are
allowed and kept apart; self-loops carry no flow. Returns (value, flow, S):
flow[i] is the flow on edge i, and S is the set of vertices that s reaches
in the final residual graph, the source side of a minimum cut.
Time O(n^3 + n m).
"""
if s == t:
raise ValueError("s and t must differ")
head: list[int] = []
cap: list[float] = []
adj: list[list[int]] = [[] for _ in range(n)]
for u, v, c in edges:
if c < 0:
raise ValueError("capacities must be nonnegative")
c = 0 if u == v else c
adj[u].append(len(head))
head.append(v)
cap.append(c)
adj[v].append(len(head))
head.append(u)
cap.append(0)
h = [0] * n # @line L1
excess = [0] * n
h[s] = n
for r in adj[s]: # @line L2
if r % 2 == 0 and cap[r] > 0:
v, d = head[r], cap[r]
cap[r] -= d
cap[r ^ 1] += d
excess[v] += d
excess[s] -= d
queue = deque(v for v in range(n) if v not in (s, t) and excess[v] > 0) # @line L3
in_queue = [False] * n
for v in queue:
in_queue[v] = True
current = [0] * n
while queue: # @line L4
u = queue.popleft() # @line L5
in_queue[u] = False
while excess[u] > 0: # @line L6
if current[u] == len(adj[u]): # @line L7
h[u] = 1 + min(h[head[r]] for r in adj[u] if cap[r] > 0) # @line L8
current[u] = 0
continue
r = adj[u][current[u]]
w = head[r]
if cap[r] > 0 and h[u] == h[w] + 1: # @line L9
d = min(excess[u], cap[r])
cap[r] -= d
cap[r ^ 1] += d
excess[u] -= d
excess[w] += d
if w != s and w != t and not in_queue[w]:
queue.append(w)
in_queue[w] = True
else:
current[u] += 1 # @line L10
flow = [cap[2 * i + 1] for i in range(len(edges))]
side, stack = {s}, [s] # @line L11
while stack:
u = stack.pop()
for r in adj[u]:
if cap[r] > 0 and head[r] not in side:
side.add(head[r])
stack.append(head[r])
return excess[t], flow, side
test_push_relabel.py 92 lines
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/push-relabel/. The authors are listed in prov/56SN6Q.jsonl.
import itertools
import random
from fractions import Fraction
import pytest
from push_relabel import push_relabel
# s, a, b, c, d, t = 0 .. 5
EXAMPLE = [(0, 1, 10), (0, 2, 10), (1, 2, 2), (1, 3, 4), (1, 4, 8),
(2, 4, 9), (4, 3, 6), (3, 5, 10), (4, 5, 10)]
def min_cut_brute(n, edges, s, t):
"""Independent oracle: the least capacity of a set S with s in S, t not in S."""
others = [v for v in range(n) if v not in (s, t)]
best = None
for r in range(len(others) + 1):
for extra in itertools.combinations(others, r):
side = {s, *extra}
c = sum(w for u, v, w in edges if u in side and v not in side)
best = c if best is None else min(best, c)
return best
def check(n, edges, s, t, value, flow, side):
assert all(0 <= f <= c for f, (_, _, c) in zip(flow, edges))
bal = [0] * n
for f, (u, v, _) in zip(flow, edges):
bal[u] -= f
bal[v] += f
assert all(bal[v] == 0 for v in range(n) if v not in (s, t))
assert bal[t] == value and -bal[s] == value
assert s in side and t not in side
assert value == sum(c for u, v, c in edges if u in side and v not in side)
def test_example() -> None:
value, flow, side = push_relabel(6, EXAMPLE, 0, 5)
assert value == 19 and side == {0, 2}
check(6, EXAMPLE, 0, 5, value, flow, side)
def test_small_example_and_stuck_excess() -> None:
edges = [(0, 1, 3), (0, 2, 2), (1, 2, 1), (1, 3, 2), (2, 3, 3)]
assert push_relabel(4, edges, 0, 3)[0] == 5
value, flow, side = push_relabel(3, [(0, 1, 5), (1, 2, 2)], 0, 2)
assert (value, flow, side) == (2, [2, 2], {0, 1})
def test_zero_capacity_and_no_path() -> None:
assert push_relabel(3, [(0, 1, 5), (1, 2, 0), (0, 2, 0)], 0, 2) == (0, [0, 0, 0], {0, 1})
assert push_relabel(4, [(0, 1, 3), (2, 3, 3)], 0, 3) == (0, [0, 0], {0, 1})
assert push_relabel(2, [], 0, 1) == (0, [], {0})
def test_antiparallel_parallel_loops() -> None:
edges = [(0, 1, 3), (1, 0, 2), (1, 2, 4), (0, 1, 2), (1, 1, 7), (2, 1, 1), (2, 0, 6)]
value, flow, side = push_relabel(3, edges, 0, 2)
assert value == 4
check(3, edges, 0, 2, value, flow, side)
def test_real_capacities() -> None:
edges = [(0, 1, Fraction(1, 2)), (0, 2, Fraction(1, 3)), (1, 2, Fraction(1, 7)),
(1, 3, Fraction(2, 5)), (2, 3, Fraction(3, 4))]
value, flow, side = push_relabel(4, edges, 0, 3)
assert value == min_cut_brute(4, edges, 0, 3) == Fraction(5, 6)
check(4, edges, 0, 3, value, flow, side)
def test_errors() -> None:
with pytest.raises(ValueError):
push_relabel(2, [(0, 1, 1)], 1, 1)
with pytest.raises(ValueError):
push_relabel(2, [(0, 1, -1)], 0, 1)
def test_random_against_brute_force() -> None:
rng = random.Random(99)
for _ in range(700):
n = rng.randint(2, 7)
m = rng.randint(0, 18)
edges = [(rng.randrange(n), rng.randrange(n), rng.choice([0, 1, 2, 3, 5, 9, 30]))
for _ in range(m)]
s, t = rng.sample(range(n), 2)
value, flow, side = push_relabel(n, edges, s, t)
assert value == min_cut_brute(n, edges, s, t)
check(n, edges, s, t, value, flow, side)
What the next level needs
- the tests did not run in this build