Clique Connect: minimum spanning tree

Question

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 S_i={A_i,1, A_i,2, ,…,A_i,Ki} consisting of K_i vertices. For every pair u,v such that u,v ∈ S_i and u<v, add an edge between vertices u and v with weight C_i.

​ 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:

The code runs okay. Ideone

from collections import defaultdict
from heapq import heappush, heappop


def solution(A):
    def prim(G):
        vis = set()
        start = next(iter(G))
        vis.add(start)
        Q, mst = [], []
        for w, nei in G[start]:
            heappush(Q, (w, start, nei))
        while len(vis) < len(G):
            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]:
                        graph[edges[jj]].append((c, edges[ii]))
                        graph[edges[ii]].append((c, edges[jj]))

    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)

Question

• How do we optimize it so that we won't get TLE? The running time must be below 2 seconds.

Answer 1

Your code with the example dataset has only 645 calls (that's very fast). Anyway, I improved it and added some comment. I don't think any particular gimmick are needed, such as using JIT.

There are some negative aspects like the names of variables, the definition of a function inside another and unused variables. I suggest you read PEP 8 — the Style Guide for Python Code.

Here the new code:

from collections import defaultdict
from heapq import heappush, heappop
from typing import List, Tuple


def prim_minimum_spanning_tree(graph: dict) -> List[Tuple[int, int, int]]:
    """
    Computes the Minimum Spanning Tree (MST) using Prim's algorithm.
    :param graph: A dictionary representing the undirected graph with edge weights.
    :return: A list of tuples (src, dest, weight) representing the MST.
    """
    visited = set()
    start_node = next(iter(graph))
    visited.add(start_node)
    min_heap, mst = [], []

    for weight, neighbor in graph[start_node]:
        heappush(min_heap, (weight, start_node, neighbor))

    while len(visited) < len(graph):
        weight, src, dest = heappop(min_heap)
        if dest in visited:
            continue
        visited.add(dest)
        mst.append((src, dest, weight))
        for weight, neighbor in graph[dest]:
            heappush(min_heap, (weight, dest, neighbor))

    return mst


def solution(A: List[List[int]]) -> int:
    """
    Computes the sum of weights of the MST for the given graph.
    :param A: A list of lists representing the graph with edge weights.
    :return: The integer solutions
    """
    graph = defaultdict(list)
    num_nodes, _ = A[0]

    for i in range(1, len(A)):
        if i % 2 == 1:
            c = A[i][1]
        else:
            edges = A[i]
            for ii in range(len(edges)):
                for jj in range(ii + 1, len(edges)):
                    if edges[ii] < edges[jj]:
                        graph[edges[jj]].append((c, edges[ii]))
                        graph[edges[ii]].append((c, edges[jj]))

    mst = prim_minimum_spanning_tree(graph)
    total_weight = sum(weight for _, _, weight in mst)

    if sorted(set(node for edge in mst for node in edge[:2])) != list(
        range(1, num_nodes + 1)
    ):
        return -1
    else:
        return total_weight


def main():
    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],
    ]
    print(solution(A))


if __name__ == "__main__":
    main()