Clique Connect: árvore geradora mínima (Kruskal vs. Prim)

Declaração do problema

Você recebe um gráfico não direcionado ponderado G com N vértices, numerados de 1 a N. Inicialmente, G não possui arestas.

Você executará M operações para adicionar arestas a G. A i-ésima operação (1≤i≤M) é a seguinte:

Você recebe um subconjunto de vértices S _i ={A _i,1 , A _i,2 , ,…,A _i,Ki } consistindo em K _i vértices. Para cada par u,v tal que u,v ∈ S _i e u<v, adicione uma aresta entre os vértices u e v com peso C _i . Depois de realizar todas as operações M, determine se G está conectado. Se for, encontre o peso total das arestas em uma árvore geradora mínima de G.

Código

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)

Pergunta

Ele produz: 4302960910. A saída esperada é 1202115217. O que eu perdi?

Alternativa:

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)

O algoritmo alternativo parece estar funcionando "parcialmente", mas ainda não tenho certeza, se foi implementado corretamente?

Pseudo-código

algorithm Kruskal(G) is
    F:= ∅
    for each v in G.V do
        MAKE-SET(v)
    for each {u, v} in G.E ordered by weight({u, v}), increasing do
        if FIND-SET(u) ≠ FIND-SET(v) then
            F := F ∪ { {u, v} }
            UNION(FIND-SET(u), FIND-SET(v))
    return F

Como executar o código online?

• Use meu modelo Ideone .
• Mude o solutionmétodo pelo método correto solution.
• Envie seu código neste link (é necessário um nome de usuário e senha de login).