I need help on coding the O(V^2) time complexity prim MST algorithm in Python. I successfully coded the O(E log E) version of the prim on this question (USACO) but due to the edge constraints the code TLEs and requires the O(V^2) version. Here is my code for the O(E log E) version:
import heapq import sys sys.stdin = open('superbull.in', 'r') sys.stdout = open('superbull.out', 'w') n = int(input()) nodes =  ids = dict() for i in range(n): nodes.append(i) ids[i] = int(input()) graph = [ for _ in range(n)] for i in range(len(nodes)): for j in range(i+1, len(nodes)): a, b = nodes[i], nodes[j] weight = ids[a]^ids[b] graph[a].append((weight, b)) graph[b].append((weight, a)) looped = [False for _ in range(n)] pq = [(0, 0)] cw, c = 0, 0 while pq: w, u = heapq.heappop(pq) w = abs(w) if not looped[u]: looped[u] = True cw += w c += 1 for nw, nn in graph[u]: if not looped[nn]: heapq.heappush(pq, (-nw, nn)) print(cw)
I attempted to code the O(V^2) version myself in Python for this question (CSES - Road Reparation):
n, m = map(int, input().split()) graph = [[float('inf') for _ in range(n)] for _ in range(n)] for _ in range(m): a, b, c = map(int, input().split()) graph[a-1][b-1] = c graph[b-1][a-1] = c print(graph) # mws =  for i in range(n): mws.append(sorted(graph[i])) looped = [False for _ in range(n)] cw, c = 0, 0 cur =  while c != n+1: for nodes in cur: mw = mws[nodes-1] for weight in mw: ...
The code isn’t complete because I got stuck addressing the following issue:
- The algorithm is that I keep track of the lightest edge for each node, however when that lightest edge forms a loop or has already been used, I do not know how to efficiently address that without making the time complexity slower than O(V^2).
- Also, I would like to know if my current O(V^2) version of the code is correct or not.