MinimalSpanningTree.py_最小生成树_源码

import sys class Graph: def __init__(self,n): self.num = n self.node = [] self.neighbors = {} self.weight = {} self.edges = [] def add_edge(self,line): s = line[0] self.node.append(s) t_list = [] num = int(line[1]) for i in range(num): end = line[2 * i + 2] w = int(line[2 * i + 3]) t_list.append(end) self.edges.append((s, end)) self.weight[(s, end)] = w self.neighbors[s] = t_list def add_in_neighbor(self,neigh, u, v): if u in neigh: neigh[u].append(v) else: neigh[u] = [v] return neigh def get_bidireact_neighbor(self,edges): neigh = {} for e in edges: s = e[0] t = e[1] neigh = self.add_in_neighbor(neigh, s, t) neigh = self.add_in_neighbor(neigh, t, s) return neigh def get_path(self,edges, u): # 判断现有边集合中两点是否可达，从而判断加入该边是否会构成圈 neigh = self.get_bidireact_neighbor(edges) if u in neigh: visit = [] path_list = [] visit += neigh[u] path_list += neigh[u] while visit: for v in visit: if v in neigh: for i in neigh[v]: if i not in path_list: visit.append(i) path_list.append(i) visit.remove(v) return path_list else: return [] def Kruskal(self): tree_edge = [] total_dist = 0 ascending_weight = sorted(self.weight.items(), key=lambda x: x[1]) for edge in ascending_weight: s = edge[0][0] t = edge[0][1] if s in self.get_path(tree_edge, t): continue else: tree_edge.append(edge[0]) total_dist += edge[1] return total_dist def Reverse_Delete(self): tree_edge = self.edges total_dist = 0 decending_weight = sorted(self.weight.items(), key=lambda x: x[1],reverse=True) for edge in decending_weight: e = edge[0] s = edge[0][0] t = edge[0][1] tmp_edge = tree_edge.copy() tmp_edge.remove(e) if s in self.get_path(tmp_edge, t): tree_edge.remove(e) for e in tree_edge: total_dist += self.weight[e] return total_dist def Prim(self): s = self.node[0] bineighbors = self.get_bidireact_neighbor(self.edges) visit = [s] tree_edge = [] total_dist = 0 while len(visit) < self.num: min_dist = float('inf') for u in visit: for v in bineighbors[u]: if v not in visit: if (u,v) in self.weight: dist = self.weight[(u,v)] else: dist = self.weight[(v,u)] if dist < min_dist: visit.append(v) min_dist = dist tree_edge.append((u,v)) total_dist += dist return total_dist if __name__ == "__main__": while True: n = int(input()) G = Graph(n) if n == 0: break else: for i in range(n-1): line = sys.stdin.readline().strip('\n').split(' ') G.add_edge(line) print('Kruskal:',G.Kruskal()) print('Reverse_Delete:',G.Reverse_Delete()) print('Prim:',G.Prim())