nauty in Python

# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/nauty/. The authors are listed in prov/RGG08G.jsonl.
"""A small model of nauty: canonical labeling by individualization-refinement.

Graphs are simple and undirected, given as (n, edges) with vertices 0..n-1.
Optional vertex colors (integers) must be kept by isomorphisms. The search
is iterative (explicit stack), prunes with automorphisms that it discovers,
and returns a canonical form, a canonical labeling, generators of the
automorphism group, and the group order.

This is a teaching model of the method in McKay's nauty, not of the
program itself: nauty uses a better refinement, a better target cell rule,
node invariants and many other speedups.
"""

Edges = list[tuple[int, int]]
Partition = list[list[int]]


def refine(adj: list[set[int]], cells: Partition) -> Partition:
    """Refine an ordered partition to the coarsest equitable one below it.

    Each pass uses every cell in order as a splitter and splits every cell
    by the number of neighbors in the splitter, pieces in increasing order
    of that number. The result does not depend on vertex names.
    Cost: O(n) passes of O(n * (n + m)) work each, so O(n^2 (n + m)).
    """
    cells = [list(c) for c in cells]
    changed = True
    while changed:
        changed = False
        for s in range(len(cells)):
            splitter = set(cells[s])
            out: Partition = []
            for cell in cells:
                if len(cell) == 1:
                    out.append(cell)
                    continue
                groups: dict[int, list[int]] = {}
                for v in cell:
                    groups.setdefault(len(adj[v] & splitter), []).append(v)
                for key in sorted(groups):
                    out.append(groups[key])
            if len(out) != len(cells):
                cells = out
                changed = True
                break
    return cells


def individualize(cells: Partition, w: int) -> Partition:
    """Split w off its cell, placing the singleton first. Cost O(n)."""
    out: Partition = []
    for cell in cells:
        if w in cell:
            out.append([w])
            rest = [v for v in cell if v != w]
            if rest:
                out.append(rest)
        else:
            out.append(cell)
    return out


def target_cell(cells: Partition) -> int:
    """Index of the first smallest cell with at least 2 vertices. Cost O(n)."""
    best = -1
    for i, cell in enumerate(cells):
        if len(cell) > 1 and (best < 0 or len(cell) < len(cells[best])):
            best = i
    return best


def _orbit_ids(n: int, gens: list[list[int]]) -> list[int]:
    """Union-find representatives of the orbits of the group generated by gens."""
    parent = list(range(n))

    def find(x: int) -> int:
        while parent[x] != x:
            parent[x] = parent[parent[x]]
            x = parent[x]
        return x

    for g in gens:
        for v in range(n):
            a, b = find(v), find(g[v])
            if a != b:
                parent[a] = b
    return [find(v) for v in range(n)]


def canonical_form(n: int, edges: Edges, colors: list[int] | None = None
                   ) -> tuple[tuple, list[int], list[list[int]], int]:
    """Return (form, labeling, generators, |Aut|).

    form is the colors and adjacency rows of the graph relabeled by labeling,
    where labeling[v] is the new name of v. Two colored graphs are isomorphic
    exactly when their forms are equal. Cost: exponential in the worst case;
    each search node costs O(n^2 (n + m)).
    """
    adj = [set() for _ in range(n)]
    for u, v in edges:
        adj[u].add(v)
        adj[v].add(u)
    colors = colors if colors is not None else [0] * n
    if n == 0:
        return ((), ()), [], [], 1

    def leaf_form(cells: Partition) -> tuple[tuple, list[int]]:
        lab = [0] * n
        for i, cell in enumerate(cells):
            lab[cell[0]] = i
        order = [cell[0] for cell in cells]
        rows = tuple(sum(1 << lab[x] for x in adj[v]) for v in order)
        return (tuple(colors[v] for v in order), rows), lab

    # @line L1
    start = refine(adj, [[v for v in range(n) if colors[v] == c] for c in sorted(set(colors))])
    gens: list[list[int]] = []
    first: tuple[tuple, list[int]] | None = None
    best: tuple[tuple, list[int]] | None = None
    first_path: list[tuple[list[int], list[int], int]] = []
    # @line L2
    # a frame: [cells, prefix, target cell, next index, tried vertices]
    stack: list[list] = [[start, [], -1, 0, []]]
    # @line L3
    while stack:
        frame = stack[-1]
        cells, prefix = frame[0], frame[1]
        # @line L4
        if len(cells) == n:
            stack.pop()
            form, lab = leaf_form(cells)
            # @line L5
            if first is None:
                first = best = (form, lab)
                continue
            for ref in (first, best):
                # @line L6
                if form == ref[0]:
                    inv = [0] * n
                    for v in range(n):
                        inv[lab[v]] = v
                    gens.append([inv[ref[1][v]] for v in range(n)])
                    break
            else:
                # @line L7
                if form > best[0]:
                    best = (form, lab)
            continue
        # @line L8
        if frame[2] < 0:
            frame[2] = target_cell(cells)
        cell = cells[frame[2]]
        on_first_path = first is None
        # @line L9
        fixing = [g for g in gens if all(g[p] == p for p in prefix)]
        orbit = _orbit_ids(n, fixing)
        tried_orbits = {orbit[t] for t in frame[4]}
        w = -1
        while frame[3] < len(cell):
            cand = cell[frame[3]]
            frame[3] += 1
            if orbit[cand] not in tried_orbits:
                w = cand
                break
        if w < 0:
            stack.pop()
            continue
        if on_first_path:
            first_path.append((list(prefix), list(cell), w))
        frame[4].append(w)
        # @line L10
        stack.append([refine(adj, individualize(cells, w)), prefix + [w], -1, 0, []])
    # @line L11
    order = 1
    for prefix, cell, w in first_path:
        fixing = [g for g in gens if all(g[p] == p for p in prefix)]
        orbit = _orbit_ids(n, fixing)
        order *= sum(1 for v in cell if orbit[v] == orbit[w])
    assert best is not None
    return best[0], best[1], gens, order


def isomorphic(g: tuple[int, Edges], h: tuple[int, Edges],
               cg: list[int] | None = None, ch: list[int] | None = None) -> bool:
    """Test isomorphism of two colored graphs by comparing canonical forms."""
    if g[0] != h[0] or len(g[1]) != len(h[1]):
        return False
    return canonical_form(g[0], g[1], cg)[0] == canonical_form(h[0], h[1], ch)[0]