设计城市之间的交通路线图

交通路线图
设计城市之间的交通路线图，有6个城市（上海、北京、南京、扬州、南通、苏州），已知城市间拟建设公路。两城市间的建设经济费用自己输入。要建设一个连接6个城市的交通网，使得任意两个城市之间都可以直接或间接到达，使得总的费用最少。
1.输入城市间的路线和建设费用；
2.输出公路网建设的最经济方案；
3.可以选择使用两种算法中的一种（Prim和Kruskal）。

# 图的表示，使用邻接矩阵
class Graph:
    def __init__(self, num_vertices):
        self.num_vertices = num_vertices
        self.graph = [[0] * num_vertices for _ in range(num_vertices)]

    def add_edge(self, src, dest, cost):
        self.graph[src][dest] = cost
        self.graph[dest][src] = cost

# Prim算法
def prim_algorithm(graph):
    num_vertices = graph.num_vertices
    selected = [False] * num_vertices  # 记录顶点是否被选中
    selected[0] = True  # 从第一个顶点开始
    min_cost = 0  # 最小生成树的总费用

    print("Prim算法最小生成树:")
    print("边\t\t费用")

    for _ in range(num_vertices - 1):  # 需要选择num_vertices-1条边
        min_edge_cost = float('inf')  # 初始化最小边的费用为正无穷
        min_edge_src = 0  # 初始化最小边的起点
        min_edge_dest = 0  # 初始化最小边的终点

        for i in range(num_vertices):
            if selected[i]:
                for j in range(num_vertices):
                    if not selected[j] and graph.graph[i][j] != 0 and graph.graph[i][j] < min_edge_cost:
                        min_edge_cost = graph.graph[i][j]
                        min_edge_src = i
                        min_edge_dest = j

        selected[min_edge_dest] = True
        print(f"{min_edge_src} - {min_edge_dest}\t\t{min_edge_cost}")
        min_cost += min_edge_cost

    print("总费用:", min_cost)

# Kruskal算法
class UnionFind:
    def __init__(self, num_vertices):
        self.num_vertices = num_vertices
        self.parent = list(range(num_vertices))

    def find(self, vertex):
        if self.parent[vertex] != vertex:
            self.parent[vertex] = self.find(self.parent[vertex])
        return self.parent[vertex]

    def union(self, x, y):
        x_root = self.find(x)
        y_root = self.find(y)
        self.parent[y_root] = x_root

def kruskal_algorithm(graph):
    num_vertices = graph.num_vertices
    edges = []  # 存储边的信息，每个元素为 (起点, 终点, 费用)
    min_cost = 0  # 最小生成树的总费用

    print("Kruskal算法最小生成树:")
    print("边\t\t费用")

    for i in range(num_vertices):
        for j in range(i + 1, num_vertices):
            if graph.graph[i][j] != 0:
                edges.append((i, j, graph.graph[i][j]))

    edges.sort(key=lambda x: x[2])  # 根据费用对边进行排序

    union_find = UnionFind(num_vertices)

    for edge in edges:
        src, dest, cost = edge
        if union_find.find(src) != union_find.find(dest):
            union_find.union(src, dest)
            print(f"{src} - {dest}\t\t{cost}")
            min_cost += cost

    print("总费用:", min_cost)

# 主程序
def main():
    num_cities = 6
    city_names = ['上海', '北京', '南京', '扬州', '南通', '苏州']
    road_costs = [
        [0, 10, 20, 30, 0, 0],
        [10, 0, 5, 0, 20, 0],
        [20, 5, 0, 0, 10, 15],
        [30, 0, 0, 0, 0, 20],
        [0, 20, 10, 0, 0, 15],
        [0, 0, 15, 20, 15, 0]
    ]

    graph = Graph(num_cities)

    for i in range(num_cities):
        for j in range(i + 1, num_cities):
            graph.add_edge(i, j, road_costs[i][j])

    prim_algorithm(graph)
    print()
    kruskal_algorithm(graph)

if __name__ == '__main__':
    main()