MST: Kruskal & Prim最小生成樹

import heapq


class DSU:
    def __init__(self, n):
        self.parent = list(range(n))
        self.size = [1] * n

    def find(self, x):
        while self.parent[x] != x:
            self.parent[x] = self.parent[self.parent[x]]   # 路徑減半
            x = self.parent[x]
        return x

    def union(self, a, b):
        a, b = self.find(a), self.find(b)
        if a == b:
            return False                        # 早就同一群：這條邊會成環
        if self.size[a] < self.size[b]:
            a, b = b, a
        self.parent[b] = a
        self.size[a] += self.size[b]
        return True


def kruskal(n, edges):
    """edges 是 (u, v, w)。回傳 (總權重, 選中的邊)，圖不連通時回傳 None。O(E log E)"""
    dsu, total, chosen = DSU(n), 0, []
    for w, u, v in sorted((w, u, v) for u, v, w in edges):
        if dsu.union(u, v):
            total += w
            chosen.append((u, v, w))
            if len(chosen) == n - 1:            # 收滿 V - 1 條就可以停
                break
    return (total, chosen) if len(chosen) == n - 1 else None


def prim(n, edges, start=0):
    """lazy Prim：堆積裡可能有兩端都在樹內的舊邊，取出時再丟掉。O(E log E)"""
    adj = [[] for _ in range(n)]
    for u, v, w in edges:
        adj[u].append((w, v))
        adj[v].append((w, u))
    in_tree, total, count = [False] * n, 0, 0
    heap = [(0, start)]
    while heap and count < n:
        w, u = heapq.heappop(heap)
        if in_tree[u]:
            continue
        in_tree[u] = True
        total += w                              # 注意是邊的權重，不是到起點的距離
        count += 1
        for e in adj[u]:
            if not in_tree[e[1]]:
                heapq.heappush(heap, e)
    return total if count == n else None


def clusters(points, k):
    """單一連結分群：Kruskal 做到剩下 k 群就停，等於把 MST 最長的 k - 1 條邊剪掉"""
    n = len(points)
    edges = sorted((abs(p[0] - q[0]) + abs(p[1] - q[1]), i, j)
                   for i, p in enumerate(points) for j, q in enumerate(points) if i < j)
    dsu, groups = DSU(n), n
    for _, i, j in edges:
        if groups == k:
            break
        if dsu.union(i, j):
            groups -= 1
    out = {}
    for i in range(n):
        out.setdefault(dsu.find(i), []).append(i)
    return sorted(out.values())


if __name__ == "__main__":
    A, B, C, D, E, F = range(6)
    edges = [(A, B, 3), (A, D, 2), (B, D, 4), (B, E, 1), (D, E, 5), (B, C, 7), (C, E, 6), (C, F, 4), (E, F, 8)]
    total, chosen = kruskal(6, edges)
    print(total, ["ABCDEF"[u] + "ABCDEF"[v] for u, v, _ in chosen])   # 16 ['BE', 'AD', 'AB', 'CF', 'CE']
    print(prim(6, edges))                       # 16：兩種方法總權重一定相同
    print(kruskal(4, [(0, 1, 1), (2, 3, 1)]))   # None：圖不連通，沒有生成樹
    pts = [(0, 0), (1, 0), (0, 1), (10, 10), (11, 10), (10, 11), (20, 0)]
    print(clusters(pts, 3))                     # [[0, 1, 2], [3, 4, 5], [6]]