Cod sursa(job #2536278)

Utilizator TeodorCotetCotet Teodor TeodorCotet Data 1 februarie 2020 18:29:20
Problema Arbore partial de cost minim Scor 40
Compilator py Status done
Runda Arhiva educationala Marime 1.19 kb
from queue import PriorityQueue

def read_gen(fname):
    with open(fname, 'rt') as fin:
        for line in fin:
            for val in line.split():
                yield int(val)

def mst(g, n, m):
    pq = PriorityQueue(maxsize=m)
    pq.put((0, 1, (0, 0)))
    s, cnt_edges = 0, 0
    nodes_taken, edges_taken = set(), set()

    while pq.empty() is False:
        c, node, (n1, n2) = pq.get()
        if node in nodes_taken:
            continue
        nodes_taken.add(node)
        s += c
        if not (n1 == 0 and n2 == 0):
            cnt_edges += 1
            edges_taken.add((n1, n2))
            
        for neigh, c in g[node]:
            pq.put((c, neigh, (node, neigh)))

    return s, cnt_edges, edges_taken

if __name__ == "__main__":
    it = read_gen('apm.in')

    n, m = next(it), next(it)
    g = {i: [] for i in range(1, n + 1)}

    for _ in range(m):
        x, y, c = next(it), next(it), next(it)
        g[x].append((y, c))
        g[y].append((x, c))

    s, cnt_edges, edges_taken = mst(g, n, m)
    
    with open('apm.out', 'wt') as fout:
        fout.write('{}\n{}\n'.format(s, cnt_edges))
        for edge in edges_taken:
            fout.write('{} {}\n'.format(edge[0], edge[1]))