2
\$\begingroup\$

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

    def add_link(
        self, vertex_one: Union[int, float, str],
        vertex_two: Union[int, float, str],
        weight: int
    ):
        """
        This will add a link from vertex_one to vertex_two.
        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[2])
        rank = [0] * 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 = [0]
        minimum_spanning_tree: List = []

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

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


graph = UndirectedGraph(6)
graph.add_link(0, 1, 4)
graph.add_link(0, 2, 4)
graph.add_link(1, 2, 2)
graph.add_link(1, 0, 4)
graph.add_link(2, 0, 4)
graph.add_link(2, 1, 2)
graph.add_link(2, 3, 3)
graph.add_link(2, 5, 2)
graph.add_link(2, 4, 4)
graph.add_link(3, 2, 3)
graph.add_link(3, 4, 3)
graph.add_link(4, 2, 4)
graph.add_link(4, 3, 3)
graph.add_link(5, 2, 2)
graph.add_link(5, 4, 3)
graph.print_min_spanning_tree(graph.kruskal_min_spanning_tree())
graph.print_min_spanning_tree(graph.prim_min_spanning_tree())
\$\endgroup\$
4

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Browse other questions tagged or ask your own question.