Jack goes to Rapture

  • + 0 comments

    Prim's Algorithm and heap

    def getCost(g_nodes, g_from, g_to, g_weight):
        # Uses Prim's algorithm to find minimal spanning tree starting from node 1, stop as soon as goal is is put in the tree. At each step record the weight of the longest edge. This weight is necessarily in the path from 1 to goal in the MSP, and will be the answer. If when the tree is done and goal is still not reached, it means no path exists.
        from heapq import heappush, heappop
        # adjacency matrix
        adj = [[] for _ in range(g_nodes + 1)]
        for i in range(len(g_weight)):
            adj[g_from[i]].append((g_weight[i], g_to[i]))
            adj[g_to[i]].append((g_weight[i], g_from[i]))
        # record which node is visited
        visited = [False, True] + [False]*(g_nodes-1)
        # adjacent edges to the tree so far. Use heap to find the one with lowest weight.
        adj_edge = []
        for edge in adj[1]:
            heappush(adj_edge, edge)
        maximum = float('-inf')
    
        while True:
            # pop the adjacent edge with lowest weight
            new_weight, new_node = heappop(adj_edge)
            while visited[new_node]:
                if len(adj_edge) == 0:
                    print('NO PATH EXISTS')
                    return
                new_weight, new_node = heappop(adj_edge)
            visited[new_node] = True
            maximum = max(maximum, new_weight)
            # if the new node just added is the goal, then we're done
            if new_node == g_nodes:
                print(maximum)
                return
            #update adjacent edges
            for edge in adj[new_node]:
                heappush(adj_edge, edge)
            
    #use adjacency matrix and heap