演算法-圖論-最小生成樹

Frank23發表於2024-11-02

1. Prim演算法(卡碼網 53)

# 最小生成樹prim演算法 
# 貪心地選擇距離已有集合最小的邊加入
def prim(graph) -> int:
    result = 0
    visited = [False for _ in range(len(graph))]
    # 未選點距離已選集合的最短距離
    minDist = [10001 for _ in range(len(graph))]
    
    # 初始化:加入節點1
    visited[1] = True
    for i in range(2, len(minDist)):
        minDist[i] = graph[1][i]
    
    # 還要加入n-1個頂點
    for i in range(len(minDist)-2):
        node_to_add = -1
        minVal = 10001
        for j in range(2, len(minDist)):
            if visited[j] == False and minDist[j] < minVal:
                minVal = minDist[j]
                node_to_add = j
        
        # 將最近的未訪問的節點加入集合
        visited[node_to_add] = True
        result += minVal
        
        # 更新minDist陣列
        for j in range(2, len(minDist)):
            if visited[j] == False and graph[node_to_add][j] < minDist[j]:
                minDist[j] = graph[node_to_add][j]
    
    return result


def main():
    num_node, num_edge = map(int, input().split())
    graph = [[10001 for _ in range(num_node+1)] for _ in range(num_node+1)]
    for _ in range(num_edge):
        node1, node2, val = map(int, input().split())
        graph[node1][node2] = val
        graph[node2][node1] = val
    
    
    print(prim(graph))
    
    
if __name__ == "__main__":
    main()

相關文章