Cod sursa(job #2689409)

Utilizator TeodorCotetCotet Teodor TeodorCotet Data 20 decembrie 2020 17:28:33
Problema Algoritmul lui Dijkstra Scor 50
Compilator py Status done
Runda Arhiva educationala Marime 1.94 kb

"""Solves dijkstra"""
from __future__ import (absolute_import, division,
                        print_function, unicode_literals)
from collections import namedtuple, defaultdict
from queue import PriorityQueue 
import math

def read_gen(fname):
    """
    Generator that reads integers one by one

    :param fname str: Path to file to read from
    """
    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):
    """Finds the shortest distance between start_node and all the other nodes.
    
    :param n: The number of nodes of the graph 
    :param graph: Dictionary containing the graph
    :param start_node: Starting node for the dijkstra algorithm
    :returns: Dictionary of the shortest distance from start_node to each node
    :rtype: dict
    """
    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):
            if dist[node] == math.inf:
                fout.write('{0} '.format(0))
            else:
                fout.write('{0} '.format(dist[node]))