An algorithm to partition a bipartite graph

matrices.py

from collections import defaultdict
from itertools import count
import numpy as np

__version__ = "2024.12.13"


def create_graph(a):
    """Returns a graph from a n x n array of bools

    The graph is a list of 2 n arrays of integers.
    The array G[i] is the lists of neighbors of
    vertex i. Vertices in [0, n) represent the rows
    of the initial array, Vertices in [n, 2n)
    represent the columns.
    """
    n, nc = a.shape
    assert n == nc and n > 0
    base = np.arange(n, dtype=int)
    G = [0] * (2 * n)
    for i, row in enumerate(a):
        G[i] = base[row] + n
    for i, col in enumerate(a.T, n):
        G[i] = base[col]
    return G


def step(g, k, dst):
    """Performs one step of the algorithm
    g is the graph, a list of size 2 n
    k is a partition of vertices of the graph,
    represented here by an array assigning an
    integer to each vertex of the graph.
    dst is a list of size 2n that will receive
    the new subpartition.
    """
    d = defaultdict(count().__next__)
    for i, neig in enumerate(g):
        m = tuple(sorted(k[j] for j in neig))
        dst[i] = d[(k[i], m)]


def algorithm(g):
    """Runs the partitioning algorithm until
    partitions are no longer be refined.

    Returns the partitioning of the graph.
    """
    k = [0] * len(g)
    kk = list(k)
    while True:
        step(g, k, kk)
        k, kk = kk, k
        if kk == k:
            break
    return k


A = np.array(
    [
        [1, 1, 0, 1, 1],
        [0, 1, 0, 1, 0],
        [0, 0, 1, 1, 0],
        [1, 0, 1, 0, 1],
        [0, 0, 1, 1, 1],
    ],
    dtype=bool,
)

print(A)
G = create_graph(A)
print(G)
k = algorithm(G)
print(k)