graduate_code/solutions/flow.py

"""G03: Edmonds-Karp with a residual matrix and a min-cut certificate."""
from collections import deque


def max_flow(capacity, source, sink):
    n = len(capacity)
    if not 2 <= n <= 64 or any(len(row) != n for row in capacity):
        raise ValueError("square capacity matrix with 2..64 vertices required")
    if any(type(c) is not int or c < 0 for row in capacity for c in row):
        raise ValueError("nonnegative integer capacities required")
    if (type(source) is not int or type(sink) is not int
            or not 0 <= source < n or not 0 <= sink < n or source == sink):
        raise ValueError("distinct valid source and sink required")
    if any(capacity[i][i] != 0 for i in range(n)):
        raise ValueError("self-loop capacities must be zero")
    residual = [list(row) for row in capacity]
    value = 0
    while True:
        parent = [-1] * n
        parent[source] = source
        todo = deque([source])
        while todo and parent[sink] == -1:
            u = todo.popleft()
            for v in range(n):
                if residual[u][v] > 0 and parent[v] == -1:
                    parent[v] = u
                    todo.append(v)
        if parent[sink] == -1:
            reachable = {i for i, p in enumerate(parent) if p != -1}
            return value, reachable, residual
        delta = None
        v = sink
        while v != source:
            u = parent[v]
            delta = residual[u][v] if delta is None else min(delta, residual[u][v])
            v = u
        v = sink
        while v != source:
            u = parent[v]
            residual[u][v] -= delta
            residual[v][u] += delta
            v = u
        value += delta


def cut_capacity(capacity, reachable):
    return sum(c for u, row in enumerate(capacity) for v, c in enumerate(row)
               if u in reachable and v not in reachable)


if __name__ == "__main__":
    c = [[0, 3, 2, 0], [0, 0, 1, 2], [0, 0, 0, 3], [0, 0, 0, 0]]
    value, side, _ = max_flow(c, 0, 3)
    print(value, sorted(side), cut_capacity(c, side))