forked from keon/algorithms
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathminimum_spanning_tree.py
More file actions
98 lines (73 loc) · 2.59 KB
/
minimum_spanning_tree.py
File metadata and controls
98 lines (73 loc) · 2.59 KB
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
"""
Minimum Spanning Tree (Kruskal's Algorithm)
Finds the MST of an undirected graph using Kruskal's algorithm with a
disjoint-set (union-find) data structure.
Reference: https://en.wikipedia.org/wiki/Kruskal%27s_algorithm
Complexity:
Time: O(E log E)
Space: O(V)
"""
from __future__ import annotations
class Edge:
"""An edge of an undirected weighted graph."""
def __init__(self, source: int, target: int, weight: int) -> None:
self.source = source
self.target = target
self.weight = weight
class DisjointSet:
"""Union-Find data structure with path compression and union by size."""
def __init__(self, size: int) -> None:
"""Create *size* singleton sets.
Args:
size: Number of elements.
"""
self.parent = list(range(size))
self.size = [1] * size
def merge_set(self, node1: int, node2: int) -> None:
"""Merge the sets containing *node1* and *node2*.
Args:
node1: First element.
node2: Second element.
"""
node1 = self.find_set(node1)
node2 = self.find_set(node2)
if self.size[node1] < self.size[node2]:
self.parent[node1] = node2
self.size[node2] += self.size[node1]
else:
self.parent[node2] = node1
self.size[node1] += self.size[node2]
def find_set(self, node: int) -> int:
"""Return the root representative of the set containing *node*.
Args:
node: Element to look up.
Returns:
Root representative.
"""
if self.parent[node] != node:
self.parent[node] = self.find_set(self.parent[node])
return self.parent[node]
def kruskal(vertex_count: int, edges: list[Edge], forest: DisjointSet) -> int:
"""Return the total weight of the MST computed by Kruskal's algorithm.
Args:
vertex_count: Number of vertices.
edges: List of weighted edges.
forest: Disjoint-set instance for the vertices.
Returns:
Sum of weights in the minimum spanning tree.
Examples:
>>> e = [Edge(0, 1, 1), Edge(1, 2, 2), Edge(0, 2, 3)]
>>> kruskal(3, e, DisjointSet(3))
3
"""
edges.sort(key=lambda edge: edge.weight)
mst: list[Edge] = []
for edge in edges:
set_u = forest.find_set(edge.source)
set_v = forest.find_set(edge.target)
if set_u != set_v:
forest.merge_set(set_u, set_v)
mst.append(edge)
if len(mst) == vertex_count - 1:
break
return sum(edge.weight for edge in mst)