1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27
 from sortedcontainers import SortedList def prim(n: int, edges: List[Tuple[int, int, int]]) > int: graph = defaultdict(dict) for u, v, w in edges: graph[u][v] = w start = 0 pq, ret, visited, prev = SortedList([(0, start, None)]), [], set(), {}
while pq: if len(visited) == n: break w, u, p = pq.pop(0) if u in visited: continue
visited.add(u), ret.append((p, u, w))
for v, w in graph[u].items(): if v in visited: continue if v in prev and prev[v][1] >= w: pq.remove((prev[v][1], v, prev[v][0])) prev[v] = (u, w) pq.add((w, v, u)) return ret[1:] if len(visited) == n else []
print(prim(5, [[0, 1, 2], [0, 2, 5], [1, 3, 3], [2, 4, 10], [3, 4, 3]]))
