Algorithm Search⌘K

Edmonds-Karp algorithm in Python

Ford-Fulkerson with BFS shortest augmenting paths, running in O(VE^2).

File
edmonds_karp.py, 68 lines
Level
I0 not checked yet
Toolchain
Python 3.14, pytest 9
Test framework
pytest
Run the tests
python3 -m pytest
edmonds_karp.pyRawGitHub
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/edmonds-karp/. The authors are listed in prov/2SESDD.jsonl.

from collections import deque
from fractions import Fraction
from typing import Union

Number = Union[int, Fraction, float]


def edmonds_karp(
    n: int, edges: list[tuple[int, int, Number]], s: int, t: int
) -> tuple[Number, list[Number], list[int]]:
    """Maximum s-t flow by shortest augmenting paths (breadth-first search).

    The network has the vertices 0 .. n-1 and the directed edges (u, v, cap) of the
    list, with cap >= 0 (int, Fraction or float). Parallel, antiparallel and
    zero-capacity edges and self-loops are allowed; each edge keeps its own flow.
    Requires s != t.
    Returns (value, flow, source_side): flow[i] is the flow on edges[i], and
    source_side lists, in increasing order, the vertices reachable from s in the final
    residual graph, the source side of a minimum cut.
    At most n * m / 2 augmentations for any capacities, so O(n m^2) time; O(n + m) space.
    """
    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[Number] = []
    adj: list[list[int]] = [[] for _ in range(n)]
    for u, v, cap in edges:  # @line L1
        if cap < 0:
            raise ValueError("capacities must be nonnegative")
        adj[u].append(len(head))
        head.append(v)
        res.append(cap)
        adj[v].append(len(head))
        head.append(u)
        res.append(0)
    value: Number = 0
    while True:
        parent_arc = [-1] * n  # @line L2
        parent_arc[s] = -2
        queue = deque([s])  # @line L3
        while queue and parent_arc[t] == -1:  # @line L4
            u = queue.popleft()
            for a in adj[u]:  # @line L5
                v = head[a]
                if res[a] > 0 and parent_arc[v] == -1:  # @line L6
                    parent_arc[v] = a  # @line L7
                    queue.append(v)
        if parent_arc[t] == -1:  # @line L8
            break
        delta = None  # @line L9
        v = t
        while v != s:
            a = parent_arc[v]
            delta = res[a] if delta is None else min(delta, res[a])
            v = head[a ^ 1]
        v = t
        while v != s:  # @line L10
            a = parent_arc[v]
            res[a] -= delta
            res[a ^ 1] += delta
            v = head[a ^ 1]
        value += delta
    flow = [res[2 * i + 1] for i in range(len(edges))]
    side = [v for v in range(n) if parent_arc[v] != -1]  # @line L11
    return value, flow, side
test_edmonds_karp.py 132 lines
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/edmonds-karp/. The authors are listed in prov/2SESDD.jsonl.

import math
import random
from fractions import Fraction
from itertools import combinations

import pytest

from edmonds_karp import edmonds_karp

S, A, B, C, D, T = range(6)
RUNNING = [(S, A, 10), (S, C, 10), (A, B, 4), (A, C, 2), (A, D, 8), (C, D, 9), (D, B, 6), (B, T, 10), (D, T, 10)]


def min_cut_brute(n, edges, s, t):
    """Oracle: the least capacity over all 2^(n-2) s-t cuts."""
    others = [v for v in range(n) if v not in (s, t)]
    best = None
    for k in range(len(others) + 1):
        for extra in combinations(others, k):
            side = {s, *extra}
            cap = sum(c for u, v, c in edges if u in side and v not in side)
            best = cap if best is None else min(best, cap)
    return best


def check(n, edges, s, t, value, flow, side):
    """The flow is feasible, has the value, and the side is a cut of that capacity."""
    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 s in side and t not in side
    cut = set(side)
    assert sum(c for u, v, c in edges if u in cut and v not in cut) == value


def test_running_example() -> None:
    value, flow, side = edmonds_karp(6, RUNNING, S, T)
    assert value == 19
    assert side == [S, C]
    check(6, RUNNING, S, T, value, flow, side)


def test_backward_edge_needed() -> None:
    edges = [(0, 1, 1), (0, 2, 1), (1, 2, 1), (1, 3, 1), (2, 3, 1)]
    value, flow, side = edmonds_karp(4, edges, 0, 3)
    assert value == 2
    check(4, edges, 0, 3, value, flow, side)


def test_no_path_and_zero_capacity() -> None:
    assert edmonds_karp(3, [(0, 1, 5), (2, 1, 5)], 0, 2) == (0, [0, 0], [0, 1])
    assert edmonds_karp(3, [(0, 1, 5), (1, 2, 0)], 0, 2) == (0, [0, 0], [0, 1])
    assert edmonds_karp(2, [], 0, 1) == (0, [], [0])


def test_antiparallel_parallel_and_loops() -> None:
    edges = [(0, 1, 3), (1, 2, 2), (2, 1, 1), (2, 3, 2), (0, 2, 1), (1, 1, 7), (0, 1, 2)]
    value, flow, side = edmonds_karp(4, edges, 0, 3)
    assert value == 2
    check(4, edges, 0, 3, value, flow, side)
    edges = [(0, 1, 4), (1, 0, 4), (1, 2, 3), (2, 1, 5)]
    value, flow, side = edmonds_karp(3, edges, 0, 2)
    assert value == 3
    check(3, edges, 0, 2, value, flow, side)


def test_exact_rational_capacities() -> None:
    edges = [(0, 1, Fraction(1, 3)), (0, 2, Fraction(1, 2)), (1, 3, Fraction(1, 2)), (2, 3, Fraction(1, 4)), (2, 1, 1)]
    value, flow, side = edmonds_karp(4, edges, 0, 3)
    assert value == Fraction(3, 4)
    check(4, edges, 0, 3, value, flow, side)


def test_bad_input() -> None:
    with pytest.raises(ValueError):
        edmonds_karp(2, [(0, 1, 1)], 0, 0)
    with pytest.raises(ValueError):
        edmonds_karp(2, [(0, 1, -1)], 0, 1)


def test_large_capacities() -> None:
    big = 2**62
    edges = [(0, 1, big), (1, 2, big), (0, 2, 1)]
    assert edmonds_karp(3, edges, 0, 2)[0] == big + 1


def test_random_against_brute_force_cut() -> None:
    rng = random.Random(4417)
    for _ in range(600):
        n = rng.randint(2, 6)
        edges = [(rng.randrange(n), rng.randrange(n), rng.randint(0, 6)) for _ in range(rng.randint(0, 12))]
        s, t = rng.sample(range(n), 2)
        value, flow, side = edmonds_karp(n, edges, s, t)
        assert value == min_cut_brute(n, edges, s, t)
        check(n, edges, s, t, value, flow, side)


def test_bad_case_for_plain_ford_fulkerson() -> None:
    x = 10**15
    edges = [(0, 1, x), (0, 2, x), (1, 2, 1), (1, 3, x), (2, 3, x)]
    value, flow, side = edmonds_karp(4, edges, 0, 3)
    assert value == 2 * x
    assert flow[2] == 0  # the shortest paths never use the middle edge


def test_irrational_like_float_capacities() -> None:
    r = (math.sqrt(5) - 1) / 2
    edges = [(0, 1, 1.0), (0, 2, r), (1, 2, 1.0), (2, 1, r * r), (1, 3, r), (2, 3, 1.0)]
    value, flow, side = edmonds_karp(4, edges, 0, 3)
    assert math.isclose(value, min_cut_brute(4, edges, 0, 3), rel_tol=1e-12)
    assert 0 in side and 3 not in side


def test_random_floats() -> None:
    rng = random.Random(77)
    for _ in range(500):
        n = rng.randint(2, 6)
        edges = [(rng.randrange(n), rng.randrange(n), rng.random() * 5) for _ in range(rng.randint(0, 12))]
        s, t = rng.sample(range(n), 2)
        value, flow, side = edmonds_karp(n, edges, s, t)
        assert math.isclose(value, min_cut_brute(n, edges, s, t), rel_tol=1e-9, abs_tol=1e-9)
        cut = set(side)
        assert math.isclose(value, sum(c for u, v, c in edges if u in cut and v not in cut), rel_tol=1e-9, abs_tol=1e-9)

What the next level needs

  • the tests did not run in this build

Read the full page about Edmonds-Karp algorithm