Help me in solving DAA140 problem

help
,
1

My issue

send solution

My code

import heapq

def primMST(graph, V):
    key = [float('inf')] * V
    key[0] = 0
    mstSet = [False] * V
    parent = [None] * V

    # Create a min heap
    pq = [(key[0], 0)]
    heapq.heapify(pq)

    while len(pq) > 0:
        u = heapq.heappop(pq)[1]
        mstSet[u] = True

        for v in range(V):
            if graph[u][v] > 0 and not mstSet[v] and key[v] > graph[u][v]:
                key[v] = graph[u][v]
                parent[v] = u
                heapq.heappush(pq, (key[v], v))

    # Calculate the total weight of the MST
    mst_weight = 0
    for i in range(1, V):
        mst_weight += graph[i][parent[i]]

    return mst_weight

# Example usage
graph = [
    [0, 2, 0, 6, 0],
    [2, 0, 3, 8, 5],
    [0, 3, 0, 0, 7],
    [6, 8, 0, 0, 9],
    [0, 5, 7, 9, 0]
]

V = len(graph)
mst_weight = primMST(graph, V)
print(mst_weight)

Learning course: Design and Analysis of Algorithms
Problem Link: https://www.codechef.com/learn/course/kl-daa/KLDAA2422/problems/DAA140