Uniformly random spanning tree over node groups

Think of each group as a single super-vertex. A set of new edges is valid exactly when, on those K super-vertices, it forms a spanning tree: K-1 edges, every group reached, no cycle between groups. Each tree edge between groups i and j can be realised by any of s_i * s_j node pairs, where s_i is the size of group i.

Part 1: any valid answer. Chain the groups: connect a node of group 0 to a node of group 1, group 1 to group 2, and so on. That is K-1 edges, connected, and acyclic by construction. O(K). Candidates should get here in a few minutes and say why K-1 is the minimum (K components need at least K-1 edges to merge).

Part 2: every valid answer equally likely. This is the real question. The obvious ideas are biased:

• Random edges with union-find (pick a random cross-group pair, keep it if it merges two components, repeat). This is Kruskal on random weights, and it does not give a uniform tree. With four singleton groups there are 16 trees; simulation shows the 4 stars come out about 10% more often than the 12 paths.
• Random tree on groups, then random endpoints. Picking a uniform labelled tree on the K groups and then uniform endpoints for each edge is wrong when sizes differ: a tree that puts a high degree on a large group has more node-level realisations, so it must be picked more often.
• Enumerate all answers and pick one. Correct, but the count is exponential.

Count first. The number of valid edge sets is

count = s_1 * s_2 * ... * s_K * N^(K-2), with N the total number of nodes.

This is the multipartite form of Cayley's formula. Summing over trees T on the groups, a tree contributes prod s_i^deg_T(i), and the sum of that over all trees is prod s_i * (s_1 + ... + s_K)^(K-2). The formula is also a recipe for sampling: it says a valid answer is equivalent to choosing one "own" node per group (s_i ways each) plus a sequence of K-2 nodes drawn from all N nodes (N ways each).

The bijection (a Prüfer code with node labels). Encode a valid edge set like a Prüfer sequence over groups: repeatedly remove the leaf group L with the smallest index. Its single remaining edge is (x in L, y in the neighbour group). Append y to the sequence and record x as L's own node. When two groups remain, the last edge gives both of their own nodes. Every group is removed once or is one of the last two, so every group gets exactly one own node, and the sequence has K-2 entries. Decoding reverses this exactly: map each sequence node to its group to get an ordinary Prüfer sequence, run the standard decode, and at each step connect the leaf's own node to the recorded node. The map is a bijection, so drawing the own nodes and the sequence uniformly and independently draws a uniform valid edge set.

Invariants during decoding. degree[g] is 1 plus the number of times group g still appears in the unread part of the sequence; a group is in the leaf heap exactly when its degree is 1 and it has not been removed. Each step removes one leaf and adds one edge between two different components, so no cycle is ever formed.

from __future__ import annotations

import heapq
import random
from typing import Optional

Edge = tuple[int, int]


def connect_groups(groups: list[list[int]]) -> list[Edge]:
    """Part 1: any valid answer. Chain the groups by their first nodes."""
    return [(groups[i][0], groups[i + 1][0]) for i in range(len(groups) - 1)]


def count_edge_sets(groups: list[list[int]]) -> int:
    """Number of valid answers: s_1 * ... * s_K * N^(K-2)."""
    k = len(groups)
    if k == 1:
        return 1
    n = sum(len(g) for g in groups)
    total = n ** (k - 2)
    for g in groups:
        total *= len(g)
    return total


def sample_edges(groups: list[list[int]],
                 rng: Optional[random.Random] = None) -> list[Edge]:
    """Part 2: return K-1 edges, every valid answer equally likely.

    Bijection: (one 'own' node per group) x (K-2 nodes drawn from all N)
    <-> valid edge sets. Decode it like a Prufer sequence over groups.
    """
    rng = rng or random.Random()
    k = len(groups)
    if k == 1:
        return []

    group_of: dict[int, int] = {}
    for gi, g in enumerate(groups):
        for node in g:
            group_of[node] = gi
    all_nodes = list(group_of)

    own = [rng.choice(g) for g in groups]                   # s_i choices each
    seq = [rng.choice(all_nodes) for _ in range(k - 2)]     # N choices each

    # Standard Prufer decode on group ids; the sequence entries also carry
    # which node inside the neighbor group receives the edge.
    degree = [1] * k
    for node in seq:
        degree[group_of[node]] += 1
    leaves = [gi for gi in range(k) if degree[gi] == 1]
    heapq.heapify(leaves)

    edges: list[Edge] = []
    for node in seq:
        leaf = heapq.heappop(leaves)
        parent = group_of[node]
        edges.append((own[leaf], node))        # leaf's own node -> recorded node
        degree[parent] -= 1
        if degree[parent] == 1:
            heapq.heappush(leaves, parent)

    u = heapq.heappop(leaves)
    v = heapq.heappop(leaves)
    edges.append((own[u], own[v]))
    return edges


def is_valid(groups: list[list[int]], edges: list[Edge]) -> bool:
    """K-1 cross-group edges that connect all groups (so no group cycle)."""
    k = len(groups)
    group_of = {node: gi for gi, g in enumerate(groups) for node in g}
    if len(edges) != k - 1:
        return False
    parent = list(range(k))

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

    for a, b in edges:
        if a not in group_of or b not in group_of:
            return False
        ra, rb = find(group_of[a]), find(group_of[b])
        if ra == rb:                           # same group or a cycle
            return False
        parent[ra] = rb
    return True

• connect_groups: O(K) time, O(K) output.
• count_edge_sets: O(K) multiplications on big integers (Python handles the size).
• sample_edges: O(N) to build the node-to-group map, O(K log K) for the Prüfer decode with a heap, O(N + K) space. With a linear-time Prüfer decode (pointer scan instead of a heap) it is O(N + K). Random draws: K + (K-2).
• is_valid: O(K α(K)) with union-find.

• groups = [[7]]: one group, return [] (count 1).
• groups = [[1], [2]]: only answer [(1, 2)].
• groups = [[1, 2], [3]]: answers {(1,3)} and {(2,3)}, each with probability 1/2 (count 2·1 = 2).
• groups = [[1], [2], [3]]: 3 answers (the three paths), each 1/3.
• groups = [[1], [2], [3], [4]]: 16 answers (Cayley 4^2); stars and paths equally likely. The random-union-find approach fails this.
• groups = [[1, 2], [3], [4, 5]]: count 2·1·2·5^1 = 20. Brute-force enumeration of all 2-edge cross-group sets that connect the groups also finds 20, and 80,000 samples hit every one of them 3,890 to 4,148 times (expected 4,000).
• groups = [[1, 2], [3, 4], [5], [6]]: count 2·2·1·1·6^2 = 144, matched by brute force; sampling is flat.
• Invalid outputs is_valid must reject: an edge inside one group ([[1, 2], [3]], [(1, 2)]), a repeated group pair that leaves a group out, wrong edge count.
• Scale: 100,000 groups of 3 nodes, the sample is valid and runs in well under a second.

1. Prove it is uniform / test it. Derive the count, show the encode and decode are inverses, then test: brute-force all answers for tiny inputs, check the count matches the formula, and run a chi-square test on sample frequencies. Interviewers like candidates who say how they would catch a subtly biased sampler, since unit tests on single outputs never will.
2. Some cross-group pairs are forbidden (a general graph between groups). The closed form no longer holds. Use Wilson's algorithm (loop-erased random walks) on the contracted multigraph, where the edge multiplicity between groups i and j is the number of allowed pairs, then pick the endpoint pair uniformly among the allowed ones. Expected time is the mean hitting time of the graph. Aldous-Broder (random walk until all groups are visited, keep first-entry edges) is simpler but slower. Mention the matrix-tree theorem for the count.
3. Weighted sampling (probability proportional to a product of edge weights). Wilson's algorithm works with transition probabilities proportional to weights. In the complete case the Prüfer trick still works if the sequence entries and own nodes are drawn with probability proportional to node weights.
4. Deterministic but balanced (minimise the maximum number of new edges on any node, or on any group). Connect in a star-free way: always attach to the node that currently has the fewest new edges, or build a balanced tree over groups; a path between groups gives max degree 2 per group and 1 per node if each group has at least two nodes.
5. Deterministic indexing (unrank the r-th answer). The bijection gives every answer a mixed-radix index: own node of group i in base s_i, then K-2 digits in base N. Unranking r in [0, count) is reading those digits and running the same decode, so you can page through answers, sample without replacement, or reproduce a sample from a single integer instead of an RNG state.

Mid-level (L4)

• Explains why K-1 edges is the minimum and why the result must be a tree over groups.
• Writes Part 1 quickly and a correct validator with union-find.
• For Part 2, proposes random edges plus union-find and, with a hint, sees that it is not uniform.

Senior (L5)

• Notices without hints that "uniform over answers" weights trees by prod s_i^deg(i), so both naive samplers are biased.
• Gets to the Prüfer-style construction or a correct rejection or Wilson approach, and implements it cleanly with the edge cases K = 1 and K = 2.
• Proposes a brute-force-plus-frequency test to check uniformity.

Staff (L6+)

• Derives prod s_i * N^(K-2) and turns the count into a bijection and an O(N + K log K) sampler, with a clear correctness argument.
• Handles the general-graph follow-up with Wilson's algorithm and the matrix-tree theorem, and states the costs.
• Discusses how to validate randomness in production code (seeded RNG injection, statistical tests in CI with fixed tolerances).