Problem Statement

You are given a weighted undirected graph G with N vertices, numbered 1 to N. Initially, G has no edges.

You will perform M operations to add edges to G. The i-th operation (1≤i≤M) is as follows:

You are given a subset of vertices Si={Ai,1, Ai,2, ,…,Ai,Ki} consisting of Ki vertices. For every pair u,v such that u,v ∈ Si and u<v, add an edge between vertices u and v with weight Ci.
​
After performing all M operations, determine whether G is connected. If it is, find the total weight of the edges in a minimum spanning tree of G.

Code

from collections import defaultdict


def solution(A):
    class Kruskal:
        def __init__(self, G):
            self.G = G
            self.parent = {}
            self.rank = {}
            self.make_sets()

        def make_sets(self):
            for u, v in self.G:
                if u not in self.parent:
                    self.parent[u] = u
                    self.rank[u] = 0
                if v not in self.parent:
                    self.parent[v] = v
                    self.rank[v] = 0

        def find(self, x):
            if self.parent[x] != x:
                self.parent[x] = self.find(self.parent[x])
            return self.parent[x]

        def union(self, u, v):
            su, sv = self.find(u), self.find(v)
            if su != sv:
                if self.rank[su] > self.rank[sv]:
                    self.parent[sv] = su
                else:
                    self.parent[su] = sv
                    if self.rank[su] == self.rank[sv]:
                        self.rank[sv] += 1

        def _mst(self):
            mst = []
            for edge in self.G.keys():
                u, v = edge
                if self.find(u) != self.find(v):
                    self.union(u, v)
                    mst.append((u, v, self.G[edge]))
            return mst

    N, M = A[0]
    graph = defaultdict(int)
    for i in range(1, len(A)):
        if i % 2 == 1:
            k, c = A[i]
        else:
            edges = A[i]
            for ii in range(len(edges)):
                for jj in range(ii + 1, len(edges)):
                    if edges[ii] < edges[jj]:
                        if (edges[jj], edges[ii]) not in graph or (edges[ii], edges[jj]) not in graph:
                            graph[(edges[jj], edges[ii])] = c
                            graph[(edges[ii], edges[jj])] = c
                            continue
                        if (edges[jj], edges[ii]) in graph and graph[(edges[jj], edges[ii])] > c:
                            graph[(edges[jj], edges[ii])] = c
                        if (edges[ii], edges[jj]) in graph and graph[(edges[ii], edges[jj])] > c:
                            graph[(edges[ii], edges[jj])] = c

    kruskal = Kruskal(graph)
    MST = kruskal._mst()
    res = 0
    nodes = set()
    # print(MST)
    for x, y, z in sorted(MST, key=lambda o: o[-1]):
        res += z
        nodes.update({x, y})

    if sorted(nodes) != list(range(1, N + 1)):
        print(-1)
    else:
        print(res)


A = [[10, 5], [6, 158260522], [1, 3, 6, 8, 9, 10], [10, 877914575], [1, 2, 3, 4, 5, 6, 7, 8, 9, 10],
     [4, 602436426], [2, 6, 7, 9], [6, 24979445], [2, 3, 4, 5, 8, 10], [4, 861648772], [2, 4, 8, 9]]

solution(A)

Question

It outputs: 4302960910. The expected output is 1202115217. What did I miss?

Alternative:

from collections import defaultdict
from heapq import heappush, heappop


def solution(A):
    def prim(G):
        vis = set()
        start, dest = next(iter(G))
        vis.add(start)
        Q, mst = [], []
        for (start, nei), w in G.items():
            heappush(Q, (w, start, nei))
        while Q:  # and len(vis) < len(G):
            # print(Q)
            w, src, dest = heappop(Q)
            if dest in vis:
                continue
            vis.add(dest)
            mst.append((src, dest, w))
            for w, nei in G[dest]:
                heappush(Q, (w, dest, nei))
        return mst

    N, M = A[0]
    graph = defaultdict(list)
    for i in range(1, len(A)):
        if i % 2 == 1:
            k, c = A[i]
        else:
            edges = A[i]
            for ii in range(len(edges)):
                for jj in range(ii + 1, len(edges)):
                    if edges[ii] < edges[jj]:
                        if (edges[jj], edges[ii]) not in graph:
                            graph[(edges[ii], edges[jj])] = c
                            graph[(edges[jj], edges[ii])] = c
                            continue

                        if (edges[jj], edges[ii]) in graph and graph[(edges[jj], edges[ii])] > c:
                            graph[(edges[jj], edges[ii])] = c
                        # if (edges[ii], edges[jj]) in graph and graph[(edges[ii], edges[jj])] > c:
                        #     graph[(edges[ii], edges[jj])] = c

    mst = prim(graph)
    res = 0
    s = set()
    for x, y, w in mst:
        res += w
        s.update({x, y})

    if sorted(s) != list(range(1, N + 1)):
        print(-1)
    else:
        print(res)


A = [[10, 5], [6, 158260522], [1, 3, 6, 8, 9, 10], [10, 877914575], [1, 2, 3, 4, 5, 6, 7, 8, 9, 10],
     [4, 602436426], [2, 6, 7, 9], [6, 24979445], [2, 3, 4, 5, 8, 10], [4, 861648772], [2, 4, 8, 9]]
solution(A)

The alternative algorithm seems to be “partially” working, but I’m still unsure, if it’s been implemented correctly?