# Kruskal's and Prim's algorithm (Minimum spanning tree)

I completed my class assignment, and I would like to have a critical code review of this implementation. How can I make this more concise, pythonic and readable? Also, while writing docstrings, is there a way to perfectly arrange the words systematically so that the docstring is consistent?

This code is about the Kruskal's algorithm and Prim's algorithm to find the minimum spanning tree. They both fall under a class of algorithms called greedy algorithms as they find the local optimum in the hopes of finding a global optimum.

Summary of Kruskal's algorithm

We start from the edges with the lowest weight and keep adding edges until we reach our goal.

The steps for implementing Kruskal's algorithm are as follows:

1. Sort all the edges from low weight to high.
2. Take the edge with the lowest weight and add it to the spanning tree.
3. If adding the edge created a cycle, then reject this edge.
4. Keep adding edges until we reach all vertices.

Summary of Prim's algorithm

The steps for implementing Prim's algorithm are as follows:

1. Initialize the minimum spanning tree with a vertex chosen at random.
2. Find all the edges that connect the tree to new vertices, find the minimum and add it to the tree.
3. Keep repeating step 2 until we get a minimum spanning tree.
from typing import List, Union

class UndirectedGraph:
def __init__(self, num_vertices: int):
"""Initializing the member variables of the graph instance."""
self.graph: List[List] = []
self.num_vertices: int = num_vertices

self, vertex_one: Union[int, float, str],
vertex_two: Union[int, float, str],
weight: int
):
"""
The graph representation will be in the form of an edge list.
"""
self.graph.append([vertex_one, vertex_two, weight])

def find_parent(
self, parent: List[int],
vertex: int
):
"""This function will find the parent of a vertex."""
if parent[vertex] == vertex:
return vertex
return self.find_parent(parent, parent[vertex])

def union(
self, parent: List,
rank: List,
node_one: int,
node_two: int
):
vertex_one = self.find_parent(parent, node_one)
vertex_two = self.find_parent(parent, node_two)
if rank[vertex_one] < rank[vertex_two]:
parent[vertex_one] = vertex_two
elif rank[vertex_one] > rank[vertex_two]:
parent[vertex_two] = vertex_one
else:
parent[vertex_two] = vertex_one
rank[vertex_one] += 1

def print_min_spanning_tree(self, tree: List[List]):
total_cost: int
total_cost = 0
for vertex_one, vertex_two, weight in tree:
print(f"{vertex_one} - {vertex_two}  |  {weight}")
total_cost += weight
print(f"Total cost of minimum spanning tree: {total_cost}")

def kruskal_min_spanning_tree(self) -> List[List]:
minimum_spanning_tree: List[List]
graph_sorted_by_weights: List[List]
rank: List[int]

minimum_spanning_tree = []
graph_sorted_by_weights = sorted(self.graph,
key=lambda element: element)
rank =  * self.num_vertices
parent = [num for num in range(self.num_vertices)]
for vertex_one, vertex_two, weight in graph_sorted_by_weights:
vertex_one_parent = self.find_parent(parent, vertex_one)
vertex_two_parent = self.find_parent(parent, vertex_two)
if vertex_one_parent != vertex_two_parent:
minimum_spanning_tree.append([vertex_one, vertex_two, weight])
self.union(parent, rank, vertex_one_parent, vertex_two_parent)
return minimum_spanning_tree

def prim_min_spanning_tree(self):
visited: List = 
minimum_spanning_tree: List = []

while len(visited) != len({i for i in self.graph}):
valid_edges = self.get_valid_edges(visited)
smallest_edge = min(valid_edges, key=lambda l: l)
minimum_spanning_tree.append(smallest_edge)
visited.append(smallest_edge)
return minimum_spanning_tree

def get_valid_edges(self, visited: List) -> List:
return [
edges for edges in self.graph
if edges in visited and edges not in visited
]

graph = UndirectedGraph(6)