Python/graphs/minimum_spanning_tree_kruskal.py

47 lines
1.3 KiB
Python
Raw Normal View History

2021-08-02 12:40:48 +00:00
def kruskal(
num_nodes: int, edges: list[tuple[int, int, int]]
) -> list[tuple[int, int, int]]:
"""
2021-08-02 12:40:48 +00:00
>>> kruskal(4, [(0, 1, 3), (1, 2, 5), (2, 3, 1)])
[(2, 3, 1), (0, 1, 3), (1, 2, 5)]
2021-08-02 12:40:48 +00:00
>>> kruskal(4, [(0, 1, 3), (1, 2, 5), (2, 3, 1), (0, 2, 1), (0, 3, 2)])
[(2, 3, 1), (0, 2, 1), (0, 1, 3)]
2021-08-02 12:40:48 +00:00
>>> kruskal(4, [(0, 1, 3), (1, 2, 5), (2, 3, 1), (0, 2, 1), (0, 3, 2),
... (2, 1, 1)])
[(2, 3, 1), (0, 2, 1), (2, 1, 1)]
"""
edges = sorted(edges, key=lambda edge: edge[2])
2018-10-19 12:48:28 +00:00
parent = list(range(num_nodes))
2018-10-19 12:48:28 +00:00
def find_parent(i):
2019-10-05 05:14:13 +00:00
if i != parent[i]:
parent[i] = find_parent(parent[i])
return parent[i]
2018-10-19 12:48:28 +00:00
minimum_spanning_tree_cost = 0
minimum_spanning_tree = []
2018-10-19 12:48:28 +00:00
for edge in edges:
parent_a = find_parent(edge[0])
parent_b = find_parent(edge[1])
2019-10-05 05:14:13 +00:00
if parent_a != parent_b:
minimum_spanning_tree_cost += edge[2]
2019-10-05 05:14:13 +00:00
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))
2021-08-02 12:40:48 +00:00
kruskal(num_nodes, edges)