Cod sursa(job #2596737)

Utilizator lev.tempfliTempfli Levente lev.tempfli Data 10 aprilie 2020 12:09:00
Problema Algoritmul lui Dijkstra Scor 50
Compilator py Status done
Runda Arhiva educationala Marime 1.32 kb
import math
import queue


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


class Graph:
    n = 0
    m = 0
    adj_list = []

    def __init__(self, ig):
        self.n, self.m = next(ig), next(ig)
        self.adj_list = [[] for _ in range(self.n + 1)]

        for _ in range(self.m):
            f, g, h = next(ig), next(ig), next(ig)
            self.adj_list[f].append((g, h))

    def dijkstra(self, k):
        dist = [math.inf for _ in range(self.n + 1)]
        pq = queue.PriorityQueue(self.m)

        dist[k] = 0
        pq.put((dist[k], k))

        while pq.empty() is False:
            top_d, top_v = pq.get()
            if top_d != dist[top_v]: continue
            for e_v, e_d in self.adj_list[top_v]:
                new_d = dist[top_v] + e_d
                if new_d < dist[e_v]:
                    dist[e_v] = new_d
                    pq.put((dist[e_v], e_v))
        return dist


if __name__ == "__main__":
    ig = input_gen_int("dijkstra.in");
    graph = Graph(ig)
    dist = graph.dijkstra(1)

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