Minimum Spanning Tree

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121

import heapq
from typing import Dict, Set, List, Any, Tuple


def mst(graph: Dict[Any, Dict[Any, float]]) -> List[Tuple[Any, Any, float]]:
    """
    Find the Minimum Spanning Tree using Prim's algorithm.

    Args:
        graph: Weighted adjacency list {node: {neighbor: weight, ...}}

    Returns:
        List of edges in MST as (from, to, weight) tuples

    Raises:
        ValueError: If graph is empty or disconnected
    """
    if not graph:
        raise ValueError("Graph is empty")

    mst_edges = []
    visited: Set[Any] = set()

    # Start from arbitrary vertex
    start = next(iter(graph))
    visited.add(start)

    # Min-heap: (weight, from_node, to_node)
    edges_heap = []
    for neighbor, weight in graph[start].items():
        heapq.heappush(edges_heap, (weight, start, neighbor))

    while edges_heap and len(visited) < len(graph):
        weight, frm, to = heapq.heappop(edges_heap)

        if to in visited:
            continue  # Skip - would create cycle

        # Add edge to MST
        mst_edges.append((frm, to, weight))
        visited.add(to)

        # Add new edges from the newly added vertex
        for neighbor, edge_weight in graph[to].items():
            if neighbor not in visited:
                heapq.heappush(edges_heap, (edge_weight, to, neighbor))

    if len(visited) != len(graph):
        raise ValueError("Graph is disconnected - MST does not exist")

    return mst_edges


def mst_weight(graph: Dict[Any, Dict[Any, float]]) -> float:
    """
    Calculate the total weight of the Minimum Spanning Tree.

    Args:
        graph: Weighted adjacency list

    Returns:
        Total weight of the MST
    """
    edges = mst(graph)
    return sum(weight for _, _, weight in edges)


def mst_as_graph(graph: Dict[Any, Dict[Any, float]]) -> Dict[Any, Dict[Any, float]]:
    """
    Return the MST as an adjacency list (undirected).

    Args:
        graph: Weighted adjacency list

    Returns:
        MST as weighted adjacency list
    """
    edges = mst(graph)
    result: Dict[Any, Dict[Any, float]] = {node: {} for node in graph}

    for frm, to, weight in edges:
        result[frm][to] = weight
        result[to][frm] = weight

    return result


# Example usage and demonstration
if __name__ == "__main__":
    # Example weighted undirected graph
    #
    #       A ---4--- B
    #       |  \      |
    #       2   3     1
    #       |    \    |
    #       C ---5--- D
    #
    example_graph = {
        "A": {"B": 4, "C": 2, "D": 3},
        "B": {"A": 4, "D": 1},
        "C": {"A": 2, "D": 5},
        "D": {"A": 3, "B": 1, "C": 5},
    }

    print("MST Edges (Prim):", mst(example_graph))
    print("MST Total Weight:", mst_weight(example_graph))
    print("MST as Graph:", mst_as_graph(example_graph))

    # Disconnected graph (should raise error)
    disconnected = {
        "A": {"B": 1},
        "B": {"A": 1},
        "C": {"D": 2},
        "D": {"C": 2},
    }

    try:
        mst(disconnected)
    except ValueError as e:
        print(f"Disconnected graph error: {e}")