MST — Kruskal আর Prim: greedy-এর প্রমাণিত জয়

import heapq


class DSU:                                   # union-find, lesson 12 থেকে
    def __init__(self, n):
        self.p = list(range(n))
        self.r = [0] * n

    def find(self, x):
        while self.p[x] != x:
            self.p[x] = self.p[self.p[x]]    # path halving
            x = self.p[x]
        return x

    def union(self, a, b):
        ra, rb = self.find(a), self.find(b)
        if ra == rb:
            return False
        if self.r[ra] < self.r[rb]:
            ra, rb = rb, ra
        self.p[rb] = ra
        if self.r[ra] == self.r[rb]:
            self.r[ra] += 1
        return True


def kruskal(n, edges):
    """edges: (w, u, v) → (total, [(u, v, w), ...])"""
    dsu = DSU(n)
    total, taken = 0, []
    for w, u, v in sorted(edges):
        if dsu.union(u, v):
            total += w
            taken.append((u, v, w))
            if len(taken) == n - 1:
                break
    assert len(taken) == n - 1, 'graph disconnected!'
    return total, taken


def prim(adj, n, start=0):
    """adj[u]: [(v, w), ...] → (total, [(u, v, w), ...])"""
    in_tree = [False] * n
    in_tree[start] = True
    heap = [(w, start, v) for v, w in adj[start]]
    heapq.heapify(heap)
    total, taken = 0, []
    while heap and len(taken) < n - 1:
        w, u, v = heapq.heappop(heap)
        if in_tree[v]:
            continue
        in_tree[v] = True
        total += w
        taken.append((u, v, w))
        for x, wx in adj[v]:
            if not in_tree[x]:
                heapq.heappush(heap, (wx, v, x))
    assert len(taken) == n - 1, 'graph disconnected!'
    return total, taken


def verify_cut_property(n, edges):
    """প্রতিটা non-trivial cut (S, V−S)-এর lightest crossing edge
    Kruskal-MST-তে আছে কি না — brute force, ছোট n-এর জন্য"""
    _, mst = kruskal(n, edges)
    mst_set = {(u, v) for u, v, _ in mst} | {(v, u) for u, v, _ in mst}
    best = {}                                # cut-mask → lightest crossing edge
    for mask in range(1, (1 << n) - 1):      # S আর V−S প্রতিসম না — অর্ধেকই যথেষ্ট
        if mask & 1 == 0:
            continue                         # vertex 0 সবসময় S-এ — duplicate এড়াতে
        for w, u, v in edges:
            if ((mask >> u) & 1) != ((mask >> v) & 1):
                if mask not in best or w < best[mask][0]:
                    best[mask] = (w, u, v)
    checked = 0
    for mask, (w, u, v) in best.items():
        checked += 1
        assert (u, v) in mst_set, (
            f'cut {mask:b}-এর lightest edge ({u},{v},{w}) MST-তে নেই!')
    return checked                           # যতটা cut পরীক্ষা হলো


# ================= unit cases =================
if __name__ == '__main__':
    # ১. লেসনের example graph — দুজনের উত্তর ১৪
    E = [(1,0,1),(2,0,3),(2,5,4),(3,3,1),(4,1,2),(5,2,5),(6,3,4),(7,1,5),(8,0,4)]
    # vertices: 0=A 1=B 2=C 3=D 4=E 5=F
    tk, mk = kruskal(6, E)
    adj = [[] for _ in range(6)]
    for w, u, v in E:
        adj[u].append((v, w)); adj[v].append((u, w))
    tp, mp = prim(adj, 6)
    assert tk == tp == 14, (tk, tp)
    assert {(u, v) for u, v, _ in mk} == {(u, v) for u, v, _ in mp}
    print(f'MST weight 14 ✓  edges: {sorted(mk)}')

    # ২. সমান weights — ভিন্ন shape হতে পারে, কিন্তু weight এক
    Esame = [(5,0,1),(5,1,2),(5,2,3),(5,3,0),(5,0,2)]
    ts, _ = kruskal(4, Esame)
    adjs = [[] for _ in range(4)]
    for w, u, v in Esame:
        adjs[u].append((v, w)); adjs[v].append((u, w))
    assert ts == prim(adjs, 4)[0] == 15
    print('equal-weight case: weight সমান ✓ (shape ভিন্ন হওয়া বৈধ)')

    # ৩. disconnected graph — স্পষ্ট ব্যর্থতা
    try:
        kruskal(4, [(1,0,1)])
        assert False
    except AssertionError as e:
        assert 'disconnected' in str(e)
        print('disconnected graph ধরা পড়েছে ✓')

    # ৪. single vertex ও two vertices
    assert kruskal(1, [])[0] == 0
    assert kruskal(2, [(7,0,1)])[0] == 7
    print('trivial cases ✓')

    # ৫. cut-property verifier — ৬ vertex-এর সব cut ঝালাই
    ncut = verify_cut_property(6, E)
    print(f'cut property: {ncut} টা cut-এর lightest edge সব MST-তে ✓')

নিজে বাড়ান:

১. boruvka(n, edges) যোগ করুন — প্রতিটা component একসাথে নিজের cheapest outgoing edge নেয়; parallel hardware-এ কেন এটা প্রিয়, timing-এ দেখুন। ২. Matrix-Prim (O(V²), no heap) যোগ করে dense regime-এ তিনজনের race চালান — experiment-এর টেবিল আরও এক কলাম বাড়ান। ৩. verify_cycle_property লিখুন — প্রতিটা cycle-এর heaviest edge MST-র বাইরে (cycle গুনতে DFS লাগবে — লেসন ১৩-র traversal কাজে লাগবে)।