Cod sursa(job #2533138)

Utilizator TeodorCotetCotet Teodor TeodorCotet Data 28 ianuarie 2020 19:45:33
Problema Algoritmul lui Dijkstra Scor 50
Compilator py Status done
Runda Arhiva educationala Marime 1.1 kb
import math
import heapq

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

def solve(n, m, g, start):
    dist = [math.inf for _ in range(n + 1)]
    dist[start] = 0
    pq = []
    heapq.heappush(pq, (dist[start], start))

    while len(pq) > 0:
        d, node = heapq.heappop(pq)
        if d > dist[node]: continue

        for nn, c in g[node]:
            new_d = dist[node] + c
            if new_d < dist[nn]:
                dist[nn] = new_d
                heapq.heappush(pq, (dist[nn], nn))
    return dist

if __name__ == "__main__":
    it = read_gen('dijkstra.in')
    n, m = next(it), next(it)
    g = {x: [] for x in range(1, n + 1)}
    for _ in range(m):
        x, y, c = next(it), next(it), next(it)
        g[x].append((y, c))
    dist = solve(n, m, g, 1)

    with open('dijkstra.out', 'wt') as fout:
        for i in range(2, n + 1):
            if dist[i] == math.inf:
                v = 0
            else:
                v = dist[i]
            fout.write('{} '.format(v))
        fout.write('\n')