Pagini recente » Cod sursa (job #2367635) | Cod sursa (job #3137627) | Cod sursa (job #346539) | Cod sursa (job #2609752) | Cod sursa (job #2689409)
"""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]))