Skip to content

Commit aa46639

Browse files
ruppysuppygithub-actions
and
github-actions
authored
Added Kruskal's Algorithm (more organized than the one present) (TheAlgorithms#2218)
* Added Kruskal's Algorithm * Added Type Hints * fixup! Format Python code with psf/black push * Added Type Hints V2 * Implemented suggestions + uniform naming convention * removed redundant variable (self.nodes) * updating DIRECTORY.md Co-authored-by: github-actions <${GITHUB_ACTOR}@users.noreply.github.com>
1 parent d25a926 commit aa46639

File tree

2 files changed

+110
-0
lines changed

2 files changed

+110
-0
lines changed

DIRECTORY.md

+1
Original file line numberDiff line numberDiff line change
@@ -267,6 +267,7 @@
267267
* [Karger](https://github.com/TheAlgorithms/Python/blob/master/graphs/karger.py)
268268
* [Minimum Spanning Tree Boruvka](https://github.com/TheAlgorithms/Python/blob/master/graphs/minimum_spanning_tree_boruvka.py)
269269
* [Minimum Spanning Tree Kruskal](https://github.com/TheAlgorithms/Python/blob/master/graphs/minimum_spanning_tree_kruskal.py)
270+
* [Minimum Spanning Tree Kruskal2](https://github.com/TheAlgorithms/Python/blob/master/graphs/minimum_spanning_tree_kruskal2.py)
270271
* [Minimum Spanning Tree Prims](https://github.com/TheAlgorithms/Python/blob/master/graphs/minimum_spanning_tree_prims.py)
271272
* [Multi Heuristic Astar](https://github.com/TheAlgorithms/Python/blob/master/graphs/multi_heuristic_astar.py)
272273
* [Page Rank](https://github.com/TheAlgorithms/Python/blob/master/graphs/page_rank.py)
+109
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,109 @@
1+
from __future__ import annotations
2+
3+
4+
class DisjointSetTreeNode:
5+
# Disjoint Set Node to store the parent and rank
6+
def __init__(self, key: int) -> None:
7+
self.key = key
8+
self.parent = self
9+
self.rank = 0
10+
11+
12+
class DisjointSetTree:
13+
# Disjoint Set DataStructure
14+
def __init__(self):
15+
# map from node name to the node object
16+
self.map = {}
17+
18+
def make_set(self, x: int) -> None:
19+
# create a new set with x as its member
20+
self.map[x] = DisjointSetTreeNode(x)
21+
22+
def find_set(self, x: int) -> DisjointSetTreeNode:
23+
# find the set x belongs to (with path-compression)
24+
elem_ref = self.map[x]
25+
if elem_ref != elem_ref.parent:
26+
elem_ref.parent = self.find_set(elem_ref.parent.key)
27+
return elem_ref.parent
28+
29+
def link(self, x: int, y: int) -> None:
30+
# helper function for union operation
31+
if x.rank > y.rank:
32+
y.parent = x
33+
else:
34+
x.parent = y
35+
if x.rank == y.rank:
36+
y.rank += 1
37+
38+
def union(self, x: int, y: int) -> None:
39+
# merge 2 disjoint sets
40+
self.link(self.find_set(x), self.find_set(y))
41+
42+
43+
class GraphUndirectedWeighted:
44+
def __init__(self):
45+
# connections: map from the node to the neighbouring nodes (with weights)
46+
self.connections = {}
47+
48+
def add_node(self, node: int) -> None:
49+
# add a node ONLY if its not present in the graph
50+
if node not in self.connections:
51+
self.connections[node] = {}
52+
53+
def add_edge(self, node1: int, node2: int, weight: int) -> None:
54+
# add an edge with the given weight
55+
self.add_node(node1)
56+
self.add_node(node2)
57+
self.connections[node1][node2] = weight
58+
self.connections[node2][node1] = weight
59+
60+
def kruskal(self) -> GraphUndirectedWeighted:
61+
# Kruskal's Algorithm to generate a Minimum Spanning Tree (MST) of a graph
62+
"""
63+
Details: https://en.wikipedia.org/wiki/Kruskal%27s_algorithm
64+
65+
Example:
66+
67+
>>> graph = GraphUndirectedWeighted()
68+
>>> graph.add_edge(1, 2, 1)
69+
>>> graph.add_edge(2, 3, 2)
70+
>>> graph.add_edge(3, 4, 1)
71+
>>> graph.add_edge(3, 5, 100) # Removed in MST
72+
>>> graph.add_edge(4, 5, 5)
73+
>>> assert 5 in graph.connections[3]
74+
>>> mst = graph.kruskal()
75+
>>> assert 5 not in mst.connections[3]
76+
"""
77+
78+
# getting the edges in ascending order of weights
79+
edges = []
80+
seen = set()
81+
for start in self.connections:
82+
for end in self.connections[start]:
83+
if (start, end) not in seen:
84+
seen.add((end, start))
85+
edges.append((start, end, self.connections[start][end]))
86+
edges.sort(key=lambda x: x[2])
87+
# creating the disjoint set
88+
disjoint_set = DisjointSetTree()
89+
[disjoint_set.make_set(node) for node in self.connections]
90+
# MST generation
91+
num_edges = 0
92+
index = 0
93+
graph = GraphUndirectedWeighted()
94+
while num_edges < len(self.connections) - 1:
95+
u, v, w = edges[index]
96+
index += 1
97+
parentu = disjoint_set.find_set(u)
98+
parentv = disjoint_set.find_set(v)
99+
if parentu != parentv:
100+
num_edges += 1
101+
graph.add_edge(u, v, w)
102+
disjoint_set.union(u, v)
103+
return graph
104+
105+
106+
if __name__ == "__main__":
107+
import doctest
108+
109+
doctest.testmod()

0 commit comments

Comments
 (0)