O(V^2) Time Complexity Prim MST Algorithm Help Needed

General Questions
#1

Hello!

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 = [1]
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.

Thank you!