5G网络建设 Prim Python

题目描述

现需要在某城市进行5G网络建设，已经选取N个地点设置5G基站，编号固定为1到N，

接下来需要各个基站之间使用光纤进行连接以确保基站能互联互通，不同基站之间架设光纤的成本各不相同，且有些节点之间已经存在光纤相连，

请你设计算法，计算出能联通这些基站的最小成本是多少。

注意，基站的联通具有传递性，入基站A与基站B架设了光纤基站B与基站C也架设了光纤，则基站A与基站C视为可以互相联通

输入描述

第一行输入表示基站的个数N，其中0<N<=20

第二行输入表示具备光纤直连条件的基站对的数目M,其中0 < M < N * (N - 1) / 2

从第三行开始连域输入M行数据，将式为：X Y Z P，

• 其中X Y表示基能的编号，0<X<=N， 0 < Y<=N 且x不等于y，
• Z表示在X Y之间架设光纤的成本，其中0 < Z < 100，
• P表示是否已存在光纤连接，0表示未连接1表示已连接.

输出描述

如果给定条件，可以建设成功互联与通的5G网络，则输出最小的建设成本， 如果给定条件，无法建设成功互联与通的5G网络，则输出-1

def help(istr: str) -> str:
    from collections import deque

    # parse input
    istr = istr.split("\n")
    N = int(istr[0])
    M = int(istr[1])
    # unable to reach all nodes
    if M < N - 1:
        return "-1"
    graph = {node: dict() for node in range(1, N + 1)}
    minimal_cost = 0
    # init Prim
    # init connected nodes and graph
    connected_nodes = []
    for i in range(M):
        edge = list(map(int, istr[2 + i].split()))
        graph[edge[0]][edge[1]] = edge[2]
        graph[edge[1]][edge[0]] = edge[2]
        if edge[3] == 1:
            for node in [edge[0], edge[1]]:
                if node not in connected_nodes:
                    connected_nodes.append(node)

    # detect orphan node
    def IsOrphanNode(graph: dict):
        visited = {node: None for node in graph}
        start = list(graph.keys())[0]
        visited[start] = 1
        visit_q = deque([start])
        while visit_q:
            current_node = visit_q.pop()
            for neighbor in graph[current_node]:
                if visited[neighbor] == None:
                    visited[neighbor] = 1
                    visit_q.appendleft(neighbor)
        visited_cnt = 0
        for _, v in visited.items():
            if v == 1:
                visited_cnt += 1
        if visited_cnt < len(graph):
            return 1  # unable to reach all nodes
        return 0

    if IsOrphanNode(graph):
        return "-1"
    # init default connected node
    if connected_nodes == []:
        connected_nodes.append(
            list(graph.keys())[0]
        )  # select first node as default start
    # init neighbor nodes
    neighbor_nodes = []
    for node in connected_nodes:
        for neighbor in graph[node]:
            if neighbor not in neighbor_nodes and neighbor not in connected_nodes:
                neighbor_nodes.append(neighbor)
    # already complete connection
    if len(connected_nodes) == len(graph):
        return "0"
    neighbor_len_list = []

    def MinLenNeighbor(
        graph: dict, nodes_a: list, nodes_b: list, len_list: list
    ) -> list:
        """return minimal edge from nodes_b to nodes_a"""
        for node_i in nodes_a:
            for node_j in nodes_b:
                if node_j in graph[node_i]:
                    if (node_j, node_i, graph[node_i][node_j]) not in len_list:
                        len_list.append((node_i, node_j, graph[node_i][node_j]))
        len_list.sort(key=lambda x: -x[2])
        return list(len_list.pop())

    # Prim main
    while len(connected_nodes) < len(graph):
        _, node_j, weight = MinLenNeighbor(
            graph, connected_nodes, neighbor_nodes, neighbor_len_list
        )
        if node_j in connected_nodes:
            continue  # nodes in selected edge already exist
        # add connected node
        connected_nodes.append(node_j)
        # remove neighbor node
        neighbor_nodes.remove(node_j)
        # update neighbor node
        for node_k in connected_nodes:
            for neighbor_k in graph[node_k]:
                if (
                    neighbor_k not in neighbor_nodes
                    and neighbor_k not in connected_nodes
                ):
                    neighbor_nodes.append(neighbor_k)
        # update minimal_cost
        minimal_cost += weight
    return str(minimal_cost)


assert (
    help(
        """\
3
3
1 2 3 0
1 3 1 0
2 3 5 0\
"""
    )
    == "4"
)

assert (
    help(
        """\
3
1
1 2 5 0\
"""
    )
    == "-1"
)


assert (
    help(
        """\
3
3
1 2 3 0
1 3 1 0
2 3 5 1\
"""
    )
    == "1"
)