Minimum cost to connect all the cities

minimum_cost_to_connect_all_cities.py

class Solution(object):

    def solve(self, cities_graph, n):
        # We have to find the minimum spanning tree cost for the given graph.
        # MST will give us the minimum cost to connect all the cities

        # store the nodes that are in the MST
        mst_nodes = set()
        # store the minimum explored cost for each node
        node_cost = {0: 0}
        # store the parent node for the nodes in the MST
        parent = {0: -1}

        for i in range(0, n):

            current_minimum_mode = self.get_minimum_node_from_explored_nodes(node_cost, mst_nodes, n)

            # update the explored nodes and their cost
            for v in range(0, n):
                if (
                    cities_graph[current_minimum_mode][v] != 0
                    and v not in mst_nodes
                    and cities_graph[current_minimum_mode][v] < node_cost.get(v, 999999999)
                ):
                    parent[v] = i
                    node_cost[v] = cities_graph[current_minimum_mode][v]

            mst_nodes.add(current_minimum_mode)

        cost = 0
        for i in range(1, n):
            cost += cities_graph[parent[i]][i]

        return cost

    def get_minimum_node_from_explored_nodes(self, node_cost, mst_nodes, n):
        # We can use a min heap to make it even more efficient
        # For now, we'll find the minimum cost node linearly
        min_cost = 9999999999
        min_node = None

        for i in range(0, n):
            if i not in mst_nodes and node_cost.get(i, 9999999999) < min_cost:
                min_cost = node_cost[i]
                min_node = i

        return min_node


if __name__ == "__main__":
    graph = [[0, 1, 1, 100, 0, 0],
             [1, 0, 1, 0, 0, 0],
             [1, 1, 0, 0, 0, 0],
             [100, 0, 0, 0, 2, 2],
             [0, 0, 0, 2, 0, 2],
             [0, 0, 0, 2, 2, 0]]
    ans = Solution().solve(graph, 6)
    print(ans)