Uniformly random spanning tree over node groups

You are given groups, a list of K disjoint, non-empty lists of integer node IDs. Nodes inside a group are already connected to each other; different groups are not connected.

You may add undirected edges between nodes of different groups. A set of new edges is valid if it connects all groups using the minimum number of edges, which is exactly K - 1. Equivalently, if you shrink each group to one vertex, the new edges form a spanning tree over the groups (no cycles, nothing left out).

Part 1. Implement connect_groups(groups) -> list[tuple[int, int]] that returns any valid set of edges.

Part 2. Implement sample_edges(groups, rng) -> list[tuple[int, int]] that returns a valid set chosen uniformly at random among all valid sets: every distinct set of K - 1 edges (edge (a, b) is the same as (b, a), order of edges does not matter) must have the same probability. Explain why your sampler is uniform and how you would test it.

Constraints

• 1 <= K <= 10^5, total number of nodes N <= 10^6.
• Node IDs are unique across all groups.
• If K = 1, return an empty list.

Example 1

groups = [[10, 11], [20]]

Valid answers: [(10, 20)] and [(11, 20)]. Part 1 may return either. Part 2 must return each with probability 1/2.

Example 2

groups = [[1], [2], [3], [4]]

There are 16 valid answers (every spanning tree on 4 labelled vertices): 4 stars such as [(1, 2), (1, 3), (1, 4)] and 12 paths such as [(2, 1), (1, 4), (4, 3)]. Part 2 must return each of the 16 with probability 1/16. Adding random cross-group edges until everything is connected does not achieve this.