Cod sursa(job #2682584)

Utilizator TeodorCotetCotet Teodor TeodorCotet Data 8 decembrie 2020 22:51:33
Problema Algoritmul lui Dijkstra Scor 0
Compilator py Status done
Runda Arhiva educationala Marime 1.23 kb

from collections import namedtuple, defaultdict
from queue import PriorityQueue 
import math

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, graph, start_node=1):
    pq = PriorityQueue()
    dist = defaultdict(lambda: math.inf)
    
    dist[start_node] = 0
    pq.put((dist[start_node], start_node))

    while pq.empty() is False:
        d, node = pq.get()
        if dist[node] < d:
            continue 
            
        for neighbour, edge_dist in graph[node]:
            if dist[neighbour] > dist[node] + edge_dist:
                dist[neighbour] = dist[node] + edge_dist 
                pq.put((dist[neighbour], neighbour))
    return dist 

if __name__ == '__main__':
    Edge = namedtuple('Edge', ['node', 'cost'])

    it = read_gen('dijkstra.in')
    n, m = next(it), next(it)
    graph = defaultdict(list)

    for _ in range(m):
        a, b, c = next(it), next(it), next(it)
        graph[a].append(Edge(node=b, cost=c))

    dist = solve(n, graph)
    with open('dijkstra.out', 'wt') as fout:
        for node in range(2, n + 1):
            fout.write(f'{dist[node]} ')