Python/graphs/minimum_spanning_tree_kruskal.py
Meysam c2a5033f9e
graphs/kruskal: add a test case to verify the correctness, fix styles (#2443)
* test/graphs/kruskal: adding a test case to verify the correctness of the algorithm

Fixes #2128

* grahps/kruskal: running psf/black

* graphs/kruskal: read edges in a friendlier fashion

Co-authored-by: Christian Clauss <cclauss@me.com>

* Update minimum_spanning_tree_kruskal.py

* fixup! Format Python code with psf/black push

* Update test_min_spanning_tree_kruskal.py

* updating DIRECTORY.md

Co-authored-by: Christian Clauss <cclauss@me.com>
Co-authored-by: John Law <johnlaw.po@gmail.com>
Co-authored-by: github-actions <${GITHUB_ACTOR}@users.noreply.github.com>
2020-10-08 20:21:48 +08:00

34 lines
949 B
Python

def kruskal(num_nodes, num_edges, edges):
edges = sorted(edges, key=lambda edge: edge[2])
parent = list(range(num_nodes))
def find_parent(i):
if i != parent[i]:
parent[i] = find_parent(parent[i])
return parent[i]
minimum_spanning_tree_cost = 0
minimum_spanning_tree = []
for edge in edges:
parent_a = find_parent(edge[0])
parent_b = find_parent(edge[1])
if parent_a != parent_b:
minimum_spanning_tree_cost += edge[2]
minimum_spanning_tree.append(edge)
parent[parent_a] = parent_b
return minimum_spanning_tree
if __name__ == "__main__": # pragma: no cover
num_nodes, num_edges = list(map(int, input().strip().split()))
edges = []
for _ in range(num_edges):
node1, node2, cost = [int(x) for x in input().strip().split()]
edges.append((node1, node2, cost))
kruskal(num_nodes, num_edges, edges)