minimum_spanning_tree_kruskal.py

defkruskal(
num_nodes: int, edges: list[tuple[int, int, int]]
) ->list[tuple[int, int, int]]:
"""
 >>> kruskal(4, [(0, 1, 3), (1, 2, 5), (2, 3, 1)])
 [(2, 3, 1), (0, 1, 3), (1, 2, 5)]

 >>> 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)]

 >>> 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=lambdaedge: edge[2])

parent=list(range(num_nodes))

deffind_parent(i):
ifi!=parent[i]:
parent[i] =find_parent(parent[i])
returnparent[i]

minimum_spanning_tree_cost=0
minimum_spanning_tree= []

foredgeinedges:
parent_a=find_parent(edge[0])
parent_b=find_parent(edge[1])
ifparent_a!=parent_b:
minimum_spanning_tree_cost+=edge[2]
minimum_spanning_tree.append(edge)
parent[parent_a] =parent_b

returnminimum_spanning_tree


if__name__=="__main__": # pragma: no cover
num_nodes, num_edges=list(map(int, input().strip().split()))
edges= []

for_inrange(num_edges):
node1, node2, cost= (int(x) forxininput().strip().split())
edges.append((node1, node2, cost))

kruskal(num_nodes, edges)