71 lines
1.5 KiB
Python
71 lines
1.5 KiB
Python
from heapq import heappush, heappop
|
|
import random
|
|
|
|
def to_vertex_list(graph):
|
|
result = {}
|
|
for a, b, kg in graph:
|
|
if a in result:
|
|
result[a].append((b, kg))
|
|
else:
|
|
result[a] = [(b, kg)]
|
|
if b in result:
|
|
result[b].append((a, kg))
|
|
else:
|
|
result[b] = [(a, kg)]
|
|
return result
|
|
|
|
def mst(graph):
|
|
graph = to_vertex_list(graph)
|
|
beg = random.choice(list(graph.keys()))
|
|
h = []
|
|
result = []
|
|
seen = set()
|
|
for c, kg in graph[beg]:
|
|
heappush(h, (kg, beg, c))
|
|
while h:
|
|
kg, beg, end = heappop(h)
|
|
# detect cycles
|
|
if end in seen:
|
|
continue
|
|
# use the edge
|
|
seen.add(beg)
|
|
seen.add(end)
|
|
result.append((beg, end))
|
|
for c, kg in graph[end]:
|
|
heappush(h, (kg, end, c))
|
|
return result
|
|
|
|
graphs = [
|
|
[
|
|
('A', 'B', 7),
|
|
('A', 'D', 5),
|
|
('B', 'D', 9),
|
|
('E', 'D', 15),
|
|
('F', 'D', 6),
|
|
('F', 'G', 11),
|
|
('F', 'E', 8),
|
|
('G', 'E', 9),
|
|
('C', 'E', 5),
|
|
('B', 'E', 7),
|
|
('B', 'C', 8),
|
|
],
|
|
[
|
|
('A', 'B', 4),
|
|
('A', 'C', 8),
|
|
('B', 'C', 11),
|
|
('B', 'E', 8),
|
|
('C', 'D', 7),
|
|
('C', 'F', 1),
|
|
('D', 'E', 2),
|
|
('D', 'F', 6),
|
|
('E', 'G', 7),
|
|
('E', 'H', 4),
|
|
('F', 'H', 2),
|
|
('G', 'H', 14),
|
|
('G', 'I', 9),
|
|
('H', 'I', 10),
|
|
],
|
|
]
|
|
|
|
for graph in graphs:
|
|
print(mst(graph))
|