Generic minimum spanning tree method in Python

Framework that builds a minimum spanning tree by repeatedly adding an edge that is safe, which Kruskal and Prim both instantiate.

File: generic_mst_method.py, 65 lines
Level: I0 not checked yet
Toolchain: Python 3.14, pytest 9
Test framework: pytest
Run the tests: python3 -m pytest

generic_mst_method.pyRaw GitHub

# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/generic-mst-method/. The authors are listed in prov/AA7K7S.jsonl.

from typing import Callable

# A cut chooser gets the tree label of each vertex and the labels of the
# trees that some edge leaves, and returns a set of tree labels S. The cut is
# (union of the trees in S, the rest). Any such set respects A.
CutChooser = Callable[[list[int], list[int]], set[int]]


def first_open_tree(comp: list[int], open_trees: list[int]) -> set[int]:
    """The default choice: the open tree with the smallest label."""
    return {min(open_trees)}


def generic_mst(
    n: int, edges: list[tuple[int, int, float]], choose_cut: CutChooser = first_open_tree
) -> list[int]:
    """Return the indices of the edges of a minimum spanning forest, in the order added.

    The vertices are 0 .. n - 1 and edges[i] = (u, v, w) is an undirected
    edge. Weights may tie (a light edge is then the one with the least
    (w, i)); the graph may be disconnected. choose_cut is the free choice of
    the method: it must return a set of tree labels that some edge leaves.

    Cost: at most n - 1 iterations of O(n + m) each, so O(n (n + m)) time
    besides choose_cut; O(n + m) extra space.
    """
    a: list[int] = []  # @line L1
    while True:
        comp = _labels(n, edges, a)
        open_trees = sorted({comp[x] for u, v, _ in edges if comp[u] != comp[v] for x in (u, v)})
        if not open_trees:  # @line L2
            return a  # @line L6
        s = choose_cut(comp, open_trees)  # @line L3
        light = None
        for i, (u, v, w) in enumerate(edges):  # @line L4
            if (comp[u] in s) != (comp[v] in s):
                if light is None or (w, i) < (edges[light][2], light):
                    light = i
        if light is None:
            raise ValueError("choose_cut returned a cut that no edge crosses")
        a.append(light)  # @line L5


def _labels(n: int, edges: list[tuple[int, int, float]], a: list[int]) -> list[int]:
    """Label each vertex with the smallest vertex of its tree in (V, a)."""
    adj: list[list[int]] = [[] for _ in range(n)]
    for i in a:
        u, v, _ = edges[i]
        adj[u].append(v)
        adj[v].append(u)
    comp = [-1] * n
    for s in range(n):
        if comp[s] < 0:
            comp[s] = s
            stack = [s]
            while stack:
                x = stack.pop()
                for y in adj[x]:
                    if comp[y] < 0:
                        comp[y] = s
                        stack.append(y)
    return comp

test_generic_mst_method.py 90 lines

# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/generic-mst-method/. The authors are listed in prov/AA7K7S.jsonl.

import itertools
import random

import pytest

from generic_mst_method import generic_mst

EXAMPLE = [(0, 1, 4), (0, 2, 3), (1, 2, 1), (1, 3, 2), (2, 3, 4), (2, 4, 5), (3, 4, 7), (3, 5, 6), (4, 5, 2)]


def components(n, edges, chosen):
    parent = list(range(n))

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

    for i in chosen:
        u, v, _ = edges[i]
        parent[find(u)] = find(v)
    return len({find(x) for x in range(n)})


def is_forest(n, edges, chosen):
    return components(n, edges, chosen) == n - len(chosen)


def brute_force(n, edges):
    k = n - components(n, edges, range(len(edges)))
    best = None
    for s in itertools.combinations(range(len(edges)), k):
        if is_forest(n, edges, s):
            w = sum(edges[i][2] for i in s)
            best = w if best is None or w < best else best
    return best


def test_example_prim_like():
    # the tree of vertex a (label 0) is always the open tree with the smallest label
    f = generic_mst(6, EXAMPLE)
    assert f == [1, 2, 3, 5, 8]
    assert sum(EXAMPLE[i][2] for i in f) == 13


def test_example_free_cuts():
    # the cuts of the page: {d}, {e}, {a}, {a, c}, {e, f}
    script = iter([3, 4, 0, 0, 4])

    def choose(comp, open_trees):
        return {comp[next(script)]}

    f = generic_mst(6, EXAMPLE, choose)
    assert f == [3, 8, 1, 2, 5]


def test_bad_cut_is_rejected():
    with pytest.raises(ValueError):
        generic_mst(3, [(0, 1, 1)], lambda comp, open_trees: {comp[2]})


def test_edge_cases():
    assert generic_mst(0, []) == []
    assert generic_mst(1, [(0, 0, 1)]) == []
    assert generic_mst(3, []) == []
    f = generic_mst(3, [(0, 1, 1), (1, 2, 1), (2, 0, 1)])
    assert len(f) == 2


def test_against_brute_force_with_random_cuts():
    rng = random.Random(211)
    for _ in range(600):
        n = rng.randrange(1, 7)
        m = rng.randrange(0, 9)
        edges = [(rng.randrange(n), rng.randrange(n), rng.randrange(0, 3)) for _ in range(m)]

        def choose(comp, open_trees):
            # a random union of trees; fall back to one open tree if no edge crosses it
            s = {t for t in set(comp) if rng.random() < 0.5}
            if any((comp[u] in s) != (comp[v] in s) for u, v, _ in edges):
                return s
            return {rng.choice(open_trees)}

        f = generic_mst(n, edges, choose)
        assert is_forest(n, edges, f)
        assert components(n, edges, f) == components(n, edges, range(m))
        assert sum(edges[i][2] for i in f) == brute_force(n, edges)

What the next level needs

the tests did not run in this build

Read the full page about Generic minimum spanning tree method