Minimum Spanning Trees

"""Chapter 26 — MST: Union-Find (Union-Find / Disjoint Set Union)
Topic: Union-Find with Path Compression and Union by Rank
Time: O(α(V)) per operation ≈ O(1) amortized
Space: O(V)

Description:
    Path Compression: During find, directly connect the parent of every node to the root.
    Union by Rank: Attach the tree with lower rank under the tree with higher rank.
    Using both optimizations together results in amortized complexity O(α(V)) ≈ O(1)."""


class UnionFind:
    """Union-Find with path compression + union by rank.

    Time: find/union O(α(n)) amortized
    Space: O(n)"""

    def __init__(self, n: int) -> None:
        """Initialize n elements. Each element is its own root.

        Time: O(n)
        Space: O(n)"""
        self.parent = list(range(n))
        self.rank = [0] * n
        self.size = [1] * n   #Set Size (Bonus)
        self.num_components = n

    def find(self, x: int) -> int:
        """Returns the root (representative element) of the set x belongs to.
        Apply path compression: update the parent of all nodes to the root while finding the root.

        Time: O(α(n)) amortized"""
        if self.parent[x] != x:
            self.parent[x] = self.find(self.parent[x])  #Path Compression
        return self.parent[x]

    def union(self, x: int, y: int) -> bool:
        """Merge the sets containing x and y.

        Returns:
            True  if the two sets are merged (they were different sets)
            False if they are already in the same set (cycle formation)

        Time: O(α(n)) amortized"""
        rx, ry = self.find(x), self.find(y)
        if rx == ry:
            return False   #Already in the Same Set

        #Union by Rank: Attach Lower Rank Under Higher Rank
        if self.rank[rx] < self.rank[ry]:
            rx, ry = ry, rx   #rx Always Has Higher Rank
        self.parent[ry] = rx
        self.size[rx] += self.size[ry]
        if self.rank[rx] == self.rank[ry]:
            self.rank[rx] += 1

        self.num_components -= 1
        return True

    def connected(self, x: int, y: int) -> bool:
        """Check if x and y are in the same set. Time O(α(n))."""
        return self.find(x) == self.find(y)

    def component_size(self, x: int) -> int:
        """Return the size of the set to which x belongs. Time O(α(n))."""
        return self.size[self.find(x)]


def main() -> None:
    print("=" * 55)
    print("Union-Find Basic Operation Test")
    print("=" * 55)

    uf = UnionFind(7)   #Nodes 0~6
    print(f"Initial Number of Sets: {uf.num_components}")

    #Union Operation
    ops = [(0, 1), (1, 2), (3, 4), (5, 6), (2, 3)]
    for u, v in ops:
        merged = uf.union(u, v)
        print(f"union({u}, {v}): {'Merged' if merged else 'Already in the Same Set'}"
              f"→ Number of Sets={uf.num_components}")

    print()
    print("Connected Status:")
    pairs = [(0, 4), (0, 5), (5, 6)]
    for u, v in pairs:
        print(f"  connected({u}, {v}): {uf.connected(u, v)}")

    print()
    print("Set Size:")
    for i in range(7):
        print(f"node {i}: root={uf.find(i)}, Set Size={uf.component_size(i)}")

    #── Path Compression Visualization ──
    print()
    print("=" * 55)
    print("Path Compression Example (Deep Chain → Compressed in 1 Step)")
    print("=" * 55)

    uf2 = UnionFind(6)
    #0→1→2→3→4→5 Chain
    for i in range(5):
        uf2.parent[i + 1] = i  #Intentionally Create Chain Bypassing Rank
    uf2.parent[0] = 0

    print(f"Parent Before Compression: {uf2.parent}")
    root = uf2.find(5)   #Path Compression Occurs
    print(f"find(5) = {root}")
    print(f"Parent After Compression: {uf2.parent}")
    print("(All Nodes Point Directly to Root)")

    #── Cycle Detection Application ──
    print()
    print("=" * 55)
    print("Cycle Detection Application")
    print("=" * 55)

    uf3 = UnionFind(4)
    edges = [(0, 1), (1, 2), (2, 3), (3, 0)]  #Last Edge Forms Cycle
    for u, v in edges:
        result = uf3.union(u, v)
        label = "Additional" if result else "Cycle Detected!"
        print(f"Edge ({u}-{v}): {label}")


if __name__ == "__main__":
    main()

"""Unit 26 — MST: Kruskal's Algorithm
Topic : Kruskal's MST Algorithm
Time  : O(E log E)
Space : O(V + E)

Description:
    After sorting edges in ascending order of weight,
    select edges that do not form a cycle to construct the MST.
    Use Union-Find to determine cycle existence in O(α(V)) time."""


class UnionFind:
    """Union-Find with path compression + rank union. Time O(α(n)) per operation."""

    def __init__(self, n: int) -> None:
        self.parent = list(range(n))
        self.rank = [0] * n

    def find(self, x: int) -> int:
        if self.parent[x] != x:
            self.parent[x] = self.find(self.parent[x])
        return self.parent[x]

    def union(self, x: int, y: int) -> bool:
        rx, ry = self.find(x), self.find(y)
        if rx == ry:
            return False
        if self.rank[rx] < self.rank[ry]:
            rx, ry = ry, rx
        self.parent[ry] = rx
        if self.rank[rx] == self.rank[ry]:
            self.rank[rx] += 1
        return True


def kruskal(
    n: int,
    edges: list[tuple[int, int, int]],
) -> tuple[int, list[tuple[int, int, int]]]:
    """Calculate the MST using Kruskal's algorithm.

    Args:
        n    : Number of vertices (0-indexed: 0 ~ n-1)
        edges: List of undirected edges (u, v, weight)

    Returns:
        (total_weight, mst_edges)
        total_weight : Total weight of the MST
        mst_edges    : List of edges included in the MST

    Time : O(E log E)
    Space: O(V + E)"""
    #1. Sort edges in ascending order of weight
    sorted_edges = sorted(edges, key=lambda e: e[2])

    uf = UnionFind(n)
    total = 0
    mst_edges: list[tuple[int, int, int]] = []

    #2. Check each edge
    for u, v, w in sorted_edges:
        if uf.union(u, v):          #No cycle → add to MST
            total += w
            mst_edges.append((u, v, w))
            if len(mst_edges) == n - 1:  #V-1 edges selected
                break

    return total, mst_edges


def main() -> None:
    #── Example 1: Lecture main graph ──
    print("=" * 55)
    print("Example 1: 6-vertex undirected graph")
    print("=" * 55)

    edges1 = [
        (0, 1, 4),   # 1-2
        (0, 2, 3),   # 1-3
        (1, 2, 1),   # 2-3
        (1, 3, 2),   # 2-4
        (1, 4, 3),   # 2-5
        (2, 4, 5),   # 3-5
        (3, 4, 6),   # 4-5
        (4, 5, 7),   # 5-6
        (3, 5, 4),   # 4-6
    ]
    total1, mst1 = kruskal(6, edges1)
    print(f"Total MST weight: {total1}")
    print("MST edges:")
    for u, v, w in mst1:
        print(f"{u+1} -- {v+1}  (weight {w})")

    #── Example 2: Complete graph (5 nodes) ──
    print()
    print("=" * 55)
    print("Example 2: Complete graph K5")
    print("=" * 55)
    import itertools
    import random
    random.seed(42)
    n2 = 5
    edges2 = []
    for i, j in itertools.combinations(range(n2), 2):
        w = random.randint(1, 20)
        edges2.append((i, j, w))

    print("Edge list:")
    for u, v, w in edges2:
        print(f"  {u}-{v}: {w}")

    total2, mst2 = kruskal(n2, edges2)
    print(f"Total MST weight: {total2}")
    print("MST edges:")
    for u, v, w in mst2:
        print(f"  {u}-{v}: {w}")

    #── Example 3: Disconnected (non-connected graph) ──
    print()
    print("=" * 55)
    print("Example 3: Disconnected graph — MSF (Minimum Spanning Forest)")
    print("=" * 55)
    edges3 = [(0, 1, 5), (1, 2, 3), (3, 4, 2)]  #0-1-2 and 3-4 are separated
    total3, mst3 = kruskal(5, edges3)
    print(f"Total MSF weight: {total3}")
    print(f"Number of selected edges: {len(mst3)} (less than V-1={4} → disconnected)")
    for u, v, w in mst3:
        print(f"  {u}-{v}: {w}")

    #── Complexity summary ──
    print()
    print("=" * 55)
    print("Kruskal complexity")
    print("=" * 55)
    print("Edge sorting: O(E log E)")
    print("  Union-Find: O(E α(V)) ≈ O(E)")
    print("Overall    : O(E log E)")
    print("Space      : O(V + E)")


if __name__ == "__main__":
    main()

"""Unit 26 — MST: Prim's Algorithm (Prim)
Topic: Prim's MST Algorithm using heapq
Time: O((V + E) log V)
Space: O(V + E)

Description:
    Start from an arbitrary vertex and repeatedly select the minimum weight edge adjacent to the MST.
    Use heapq (min-heap) to quickly select the minimum weight edge among unvisited vertices.
    Lazy deletion: If a vertex popped from the heap is already included in the MST, skip it."""

import heapq
import sys

INF = sys.maxsize


def prim(
    graph: list[list[tuple[int, int]]],
    start: int = 0,
) -> tuple[int, list[tuple[int, int, int]]]:
    """Calculate the MST using Prim's algorithm.

    Args:
        graph: graph[u] = [(v, weight), ...] adjacency list (undirected)
        start: starting vertex (0-indexed)

    Returns:
        (total_weight, mst_edges)

    Time: O((V + E) log V)
    Space: O(V + E)"""
    n = len(graph)
    in_mst = [False] * n
    total = 0
    mst_edges: list[tuple[int, int, int]] = []

    #(weight, current node, previous node)
    heap: list[tuple[int, int, int]] = [(0, start, -1)]

    while heap and len(mst_edges) < n - 1:
        w, u, prev = heapq.heappop(heap)

        if in_mst[u]:       #lazy deletion: already included in MST
            continue

        in_mst[u] = True
        total += w

        if prev != -1:
            mst_edges.append((prev, u, w))

        #add only unvisited adjacent vertices to the heap
        for v, weight in graph[u]:
            if not in_mst[v]:
                heapq.heappush(heap, (weight, v, u))

    return total, mst_edges


def build_undirected_graph(
    n: int,
    edges: list[tuple[int, int, int]],
) -> list[list[tuple[int, int]]]:
    """Create an undirected adjacency list. Time O(E), Space O(V+E)."""
    graph: list[list[tuple[int, int]]] = [[] for _ in range(n)]
    for u, v, w in edges:
        graph[u].append((v, w))
        graph[v].append((u, w))
    return graph


def main() -> None:
    #── Example 1: Basic Graph ──
    print("=" * 55)
    print("Example 1: Prim MST (5 vertices)")
    print("=" * 55)

    #Nodes 0~4
    # 0──4──1──2──3
    # |      \  |
    # 8       1 3
    # |          \
    # 4 ────5──── 2
    edges1 = [
        (0, 1, 4), (0, 4, 8),
        (1, 2, 2), (1, 3, 1),
        (2, 3, 3), (3, 4, 5),
    ]
    graph1 = build_undirected_graph(5, edges1)
    total1, mst1 = prim(graph1, start=0)

    print(f"Total MST weight: {total1}")
    print("MST edges:")
    for u, v, w in mst1:
        print(f"{u} -- {v}  (weight {w})")

    #── Example 2: Comparison with Kruskal ──
    print()
    print("=" * 55)
    print("Example 2: Prim vs Kruskal Result Comparison")
    print("=" * 55)

    edges2 = [
        (0, 1, 7), (0, 3, 5),
        (1, 2, 8), (1, 3, 9), (1, 4, 7),
        (2, 4, 5),
        (3, 4, 15), (3, 5, 6),
        (4, 5, 8), (4, 6, 9),
        (5, 6, 11),
    ]
    n2 = 7
    graph2 = build_undirected_graph(n2, edges2)
    total_prim, mst_prim = prim(graph2, start=0)

    #Kruskal (comparison with simple implementation)
    parent = list(range(n2))

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

    def union(x: int, y: int) -> bool:
        rx, ry = find(x), find(y)
        if rx == ry:
            return False
        parent[ry] = rx
        return True

    sorted_e = sorted(edges2, key=lambda e: e[2])
    total_kr = 0
    for u, v, w in sorted_e:
        if union(u, v):
            total_kr += w

    print(f"Prim    MST weight: {total_prim}")
    print(f"Kruskal MST weight: {total_kr}")
    print(f"Results identical: {total_prim == total_kr}")

    #── Example 3: Coordinate-based Complete Graph (Prim Advantage) ──
    print()
    print("=" * 55)
    print("Example 3: Coordinate-based Complete Graph MST")
    print("— Weights set by Euclidean distance between coordinates")
    print("=" * 55)

    import math
    coords = [(0, 0), (3, 0), (6, 1), (1, 4), (4, 5)]
    m = len(coords)
    edges3 = []
    for i in range(m):
        for j in range(i + 1, m):
            x1, y1 = coords[i]
            x2, y2 = coords[j]
            dist = math.sqrt((x1 - x2) ** 2 + (y1 - y2) ** 2)
            edges3.append((i, j, dist))

    graph3_list: list[list[tuple[int, float]]] = [[] for _ in range(m)]
    for u, v, w in edges3:
        graph3_list[u].append((v, w))  # type: ignore[arg-type]
        graph3_list[v].append((u, w))  # type: ignore[arg-type]

    #prim function expects int weight; float version inline
    in_mst = [False] * m
    total3 = 0.0
    heap3: list[tuple[float, int, int]] = [(0.0, 0, -1)]
    mst3: list[tuple[int, int, float]] = []

    while heap3 and len(mst3) < m - 1:
        w, u, prev = heapq.heappop(heap3)
        if in_mst[u]:
            continue
        in_mst[u] = True
        total3 += w
        if prev != -1:
            mst3.append((prev, u, w))
        for v, weight in graph3_list[u]:
            if not in_mst[v]:
                heapq.heappush(heap3, (weight, v, u))

    print(f"Total MST distance: {total3:.4f}")
    for u, v, w in mst3:
        print(f"{coords[u]} -- {coords[v]}  distance={w:.4f}")


if __name__ == "__main__":
    main()

"""Chapter 26 — MST: Applications (MST Weights, Second MST, Bottleneck Spanning Tree)
Topic: MST Applications — Second MST, Bottleneck Spanning Tree
Time: MST O(E log E), Second MST O(E * V)
Space: O(V + E)

Description:
    1. Calculate the total weight of the MST (basic)
    2. Second Minimum Spanning Tree
       — Remove one edge from the MST and add a replacement edge
    3. Bottleneck Spanning Tree
       — A spanning tree with the minimum possible maximum edge weight (same as MST)"""


class UnionFind:
    """Path Compression + Union by Rank. Time O(α(n)), Space O(n)."""

    def __init__(self, n: int) -> None:
        self.parent = list(range(n))
        self.rank = [0] * n

    def find(self, x: int) -> int:
        if self.parent[x] != x:
            self.parent[x] = self.find(self.parent[x])
        return self.parent[x]

    def union(self, x: int, y: int) -> bool:
        rx, ry = self.find(x), self.find(y)
        if rx == ry:
            return False
        if self.rank[rx] < self.rank[ry]:
            rx, ry = ry, rx
        self.parent[ry] = rx
        if self.rank[rx] == self.rank[ry]:
            self.rank[rx] += 1
        return True


def kruskal(
    n: int,
    edges: list[tuple[int, int, int]],
) -> tuple[int, list[tuple[int, int, int]]]:
    """Kruskal MST. Time O(E log E), Space O(V+E)."""
    sorted_edges = sorted(edges, key=lambda e: e[2])
    uf = UnionFind(n)
    total = 0
    mst_edges: list[tuple[int, int, int]] = []
    for u, v, w in sorted_edges:
        if uf.union(u, v):
            total += w
            mst_edges.append((u, v, w))
            if len(mst_edges) == n - 1:
                break
    return total, mst_edges


# ──────────────────────────────────────────────
#Second Minimum Spanning Tree
# ──────────────────────────────────────────────

def second_mst(
    n: int,
    edges: list[tuple[int, int, int]],
) -> int:
    """Returns the weight of the Second Minimum Spanning Tree.

    Algorithm:
    1. Compute the MST.
    2. Remove each edge e from the MST one by one.
    3. Compute the MST again with the remaining edges (Kruskal without e).
    4. Return the minimum weight among the valid spanning trees.

    Time: O(E * (V + E log E))  Repeating MST E times
    Space: O(V + E)

    Note: There is a more efficient O(E log E) algorithm, but it is complicated to implement."""
    _, mst_edges = kruskal(n, edges)
    mst_set = set(map(id, mst_edges))  #Use index since it is difficult to identify the same object

    #Practical implementation: track MST edge indices
    sorted_edges = sorted(edges, key=lambda e: e[2])
    uf_init = UnionFind(n)
    mst_idx: list[int] = []
    for idx, (u, v, w) in enumerate(sorted_edges):
        if uf_init.union(u, v):
            mst_idx.append(idx)
            if len(mst_idx) == n - 1:
                break

    best = float("inf")
    #Exclude one MST edge at a time and calculate a new MST
    for skip in mst_idx:
        uf2 = UnionFind(n)
        total2 = 0
        count2 = 0
        for idx, (u, v, w) in enumerate(sorted_edges):
            if idx == skip:
                continue
            if uf2.union(u, v):
                total2 += w
                count2 += 1
                if count2 == n - 1:
                    break
        if count2 == n - 1:
            best = min(best, total2)

    return int(best) if best != float("inf") else -1


# ──────────────────────────────────────────────
#Bottleneck Spanning Tree
# ──────────────────────────────────────────────

def bottleneck_mst(
    n: int,
    edges: list[tuple[int, int, int]],
) -> int:
    """Returns the maximum edge weight of the Bottleneck Spanning Tree.

    Note: The MST is also a Bottleneck Spanning Tree.
    (Among all spanning trees, the one with the minimum maximum edge = MST)

    Time : O(E log E)
    Space: O(V + E)"""
    _, mst_edges = kruskal(n, edges)
    if len(mst_edges) < n - 1:
        return -1  #Cannot connect
    return max(w for _, _, w in mst_edges)


def main() -> None:
    #── Common Example Graph ──
    n = 5
    edges = [
        (0, 1, 1), (0, 2, 3),
        (1, 2, 1), (1, 3, 4), (1, 4, 2),
        (2, 3, 5),
        (3, 4, 7),
    ]

    #1. MST Weight
    print("=" * 55)
    print("1. MST Weight")
    print("=" * 55)
    total, mst = kruskal(n, edges)
    print(f"MST Weight: {total}")
    print("MST edges:")
    for u, v, w in mst:
        print(f"  {u}--{v}: {w}")

    #2. Second MST
    print()
    print("=" * 55)
    print("2. Second Minimum Spanning Tree")
    print("=" * 55)
    second = second_mst(n, edges)
    print(f"Second MST Weight: {second}")
    print(f"(Difference from the first MST: {second - total})")

    #3. Bottleneck Spanning Tree
    print()
    print("=" * 55)
    print("3. Maximum Edge of Bottleneck Spanning Tree")
    print("=" * 55)
    bottleneck = bottleneck_mst(n, edges)
    print(f"Maximum edge weight of MST (Bottleneck): {bottleneck}")

    #Compare with another spanning tree
    #For comparison: one of the spanning trees including all edges arbitrarily
    other_st = [(0, 1, 1), (1, 2, 1), (1, 4, 2), (1, 3, 4)]
    other_max = max(w for _, _, w in other_st)
    print(f"Maximum edge of other spanning tree: {other_max}")
    print(f"MST is bottleneck-optimal: {bottleneck <= other_max}")

    #── Application Summary ──
    print()
    print("=" * 55)
    print("MST Application Pattern Summary")
    print("=" * 55)
    print("MST Weight    : kruskal/prim basic")
    print("Second MST    : recalculate after removing MST edges")
    print("Bottleneck Spanning Tree : MST = minimum bottleneck tree")
    print("Clustering    : remove k-1 largest edges → k clusters")
    print("Network Cost  : MST = minimum connection cost")


if __name__ == "__main__":
    main()