prim算法python实现

以下是Prim算法的Python实现：

INF = 9999999  # 定义一个无穷大的数

def prim(graph, start):
    """
    :param graph: 图
    :param start: 起点
    :return: 最小生成树
    """
    n = len(graph)  # 图的大小
    lowcost = [INF] * n  # 到达每个节点的最小权值
    closest = [start] * n  # 最小权值的边所连接的节点
    visited = [False] * n  # 标记节点是否被访问过
    visited[start] = True
    for i in range(n - 1):
        min_cost = INF
        for j in range(n):
            if not visited[j] and graph[closest[j]][j] < lowcost[j]:
                lowcost[j] = graph[closest[j]][j]
                closest[j] = start
            if not visited[j] and lowcost[j] < min_cost:
                min_cost = lowcost[j]
                k = j
        visited[k] = True
        for j in range(n):
            if not visited[j] and graph[k][j] < lowcost[j]:
                lowcost[j] = graph[k][j]
                closest[j] = k
    return closest

# 测试
graph = [[0, 1, 3, INF],
         [1, 0, 3, 6],
         [3, 3, 0, 4],
         [INF, 6, 4, 0]]
start = 0
print(prim(graph, start))  # 输出：[0, 0, 1, 2]