Prim、Kruskal求最小生成树【Python】

算法作业，没有考虑非连通图

1.给定一个无向图，自选一种数据存储结构，实现最小生成树计算的Prim算法和Kruskal算法，分析基于所选数据存储结构得到的算法复杂度。（100分）
输入：
第一行输入两个数m, n，表图示有m个顶点（所有顶点的字母各不相同），n条边；
接下来n行每行输入两个顶点，一个值w，表示这两个顶点之间有边连接，且边的权重为w；

输出:
分别输出Prim算法和Kruskal算法得到的最小生成树（若某个节点存在最小生成树的多个选择，则按照字母表顺序进行优先选择，即输出只有一种结果）
例子：
输入：
3,3
A,B,2
A,C,5
B,C,4
输出：
Prim：
A,B,2
B,C,4
Kruskal:
A,B,2
B,C,4

 

 代码如下：

# 选出最小字母
def sort(V):
    min = 0
    for i in range(1, len(V)):
        if ord(V[i]) < ord(V[min]):
            min = i
    return V[min]


def sort_G(G):
    n = len(G)
    for i in range(n):
        for j in range(n - i - 1):
            if G[j][2] > G[j + 1][2]:
                G[j], G[j + 1] = G[j + 1], G[j]
    return G


def PRIM(G, T, U, V, dist, min):
    # 从最小的点开始
    U.append(min)
    V.remove(min)
    dist[min] = 0
    for v1, v2, w in G:
        if v1 == min:
            temp = v2  # 相邻边
        elif v2 == min:
            temp = v1
        else:
            continue
        dist[temp] = w
    prim(G, T, U, V, dist)


def prim(G, T, U, V, dist):
    if not V:
        return

    # 需要更新的点集
    upgrade_dist = []

    change = False
    min = V[0]

    for i in range(len(V)):
        if dist[V[i]] < dist[min]:
            min = V[i]
        elif dist[V[i]] == dist[min]:
            if ord(V[i]) < ord(min):
                min = V[i]
        else:
            continue

    for v1, v2, w in G:
        if v1 == min:
            temp = v2  # 相邻边
        elif v2 == min:
            temp = v1
        else:
            continue

        # 更新距离
        if temp not in U and (dist[temp] > w):
            dist[temp] = w
            upgrade_dist.append(temp)
        elif temp in U and change == False:
            if (min, temp, dist[min]) in G:
                T.append((min, temp, dist[min]))
                V.remove(min)
                U.append(min)
                change = True
            elif (temp, min, dist[min]) in G:
                T.append((temp, min, dist[min]))
                V.remove(min)
                U.append(min)
                change = True

    prim(G, T, U, V, dist)


def exchange(V, v1, v2):
    if V[v1] < V[v2]:
        for v, k in V.items():
            if V[v] == V[v2]:
                V[v] = V[v1]
    else:
        V[v1] = V[v2]
        for v, k in V.items():
            if V[v] == V[v1]:
                V[v] = V[v2]


def Kruskal(G, V, T):

    if len(T) >= len(V) - 1:
        return

    v1, v2, w = G[0]
    if V[v1] != V[v2]:
        exchange(V, v1, v2)
        T.append((v1, v2, w))
        # 删去
        G.pop(0)
    else:
        G.pop(0)
    Kruskal(G, V, T)

if __name__ == "__main__":
    while True:
        choice = input("请输入：e:Prim + Kruskal q:退出")
        if choice == "q":
            break
        m, n = map(int, input("请输入点和边的个数，以逗号分隔:").split(","))
        print(m, n)
        print("按序输入：相邻两点，边的权重，以逗号分隔：")
        # 图的信息、生成树、生成树点集、未进入生成树的点集
        G, T, U, V = [], [], [], []

        for i in range(n):
            V_1, V_2, E = input("").split(",")
            G.append((V_1, V_2, int(E)))
            V.append(V_1)
            V.append(V_2)

        # 去重
        V = list(set(V))

        # Kruskal

        # G 按权值排序
        sorted_G = sort_G(G).copy()
        # 结点编号
        V_k = {}
        i = 0
        for v in V:
            V_k[v] = i
            i += 1

        T_k = T.copy()
        Kruskal(sorted_G, V_k, T_k)
        print("Kruskal最小生成树为：")
        for v1, v2, w in T_k:
            print(v1, v2, w, )

        # Prim

        # 距离，字典存储，后续更新
        dist = dict.fromkeys(V, float('inf'))

        PRIM(G, T, U, V, dist, sort(V))

        print("Prim最小生成树为：")
        for v1, v2, w in T:
            print(v1, v2, w, )