gomory_hu.py

"""Create a Gomory--Hu tree $T$ of an unweighted and undirected $G$.

The minimum $s$-$t$-cut in $G$ is the minimum weight edge on the unique
shortest $s$-$t$-path in $T$.

"""

import networkx as nx


def gomory_hu(G: nx.Graph) -> nx.Graph:
    V = list(G.nodes)
    T = nx.Graph()
    T.add_nodes_from(V)

    if len(V) <= 1:
        return T

    root = V[0]
    parent = {v: root for v in V if v != root}
    weight = {}

    for s in V[1:]:
        t = parent[s]

        cut_value, (A, B) = minimum_cut(G, s, t)
        A, B = set(A), set(B)

        # Ensure A is the side containing s.
        if s not in A:
            A, B = B, A

        weight[s] = cut_value

        for v in V[1:]:
            if v != s and parent[v] == t and v in A:
                parent[v] = s

        if t != root and parent[t] in A:
            parent[s] = parent[t]
            parent[t] = s
            weight[s] = weight[t]
            weight[t] = cut_value

    for v in V[1:]:
        T.add_edge(v, parent[v], weight=weight[v])

    return T