fork download
  1. from collections import defaultdict
  2. from heapq import heappush, heappop
  3.  
  4.  
  5. def solution(A):
  6. def prim(G):
  7. vis = set()
  8. start = next(iter(G))
  9. vis.add(start)
  10. Q, mst = [], []
  11. for w, nei in G[start]:
  12. heappush(Q, (w, start, nei))
  13. while len(vis) < len(G):
  14. w, src, dest = heappop(Q)
  15. if dest in vis:
  16. continue
  17. vis.add(dest)
  18. mst.append((src, dest, w))
  19. for w, nei in G[dest]:
  20. heappush(Q, (w, dest, nei))
  21. return mst
  22.  
  23. N, M = A[0]
  24. graph = defaultdict(list)
  25. for i in range(1, len(A)):
  26. if i % 2 == 1:
  27. k, c = A[i]
  28. else:
  29. edges = A[i]
  30. for ii in range(len(edges)):
  31. for jj in range(ii + 1, len(edges)):
  32. if edges[ii] < edges[jj]:
  33. graph[edges[jj]].append((c, edges[ii]))
  34. graph[edges[ii]].append((c, edges[jj]))
  35.  
  36. mst = prim(graph)
  37. res = 0
  38. s = set()
  39. for x, y, w in mst:
  40. res += w
  41. s.update({x, y})
  42.  
  43. if sorted(s) != list(range(1, N + 1)):
  44. print(-1)
  45. else:
  46. print(res)
  47.  
  48.  
  49. A = [[10, 5], [6, 158260522], [1, 3, 6, 8, 9, 10], [10, 877914575], [1, 2, 3, 4, 5, 6, 7, 8, 9, 10],
  50. [4, 602436426], [2, 6, 7, 9], [6, 24979445], [2, 3, 4, 5, 8, 10], [4, 861648772], [2, 4, 8, 9]]
  51. solution(A)
  52.  
Success #stdin #stdout 0.03s 9856KB
stdin
Standard input is empty
stdout
1202115217