Cod sursa(job #2307216)

Utilizator 24601Dan Ban 24601 Data 24 decembrie 2018 00:06:21
Problema Algoritmul lui Dijkstra Scor 50
Compilator py Status done
Runda Arhiva educationala Marime 2.52 kb
import heapq

class MyHeap:
    def __init__(self, nodes):
        self.data = [(0, 0), ] * (nodes + 1)
        self.nodeToHeap = {}
        self.size = 0

    def insert(self, n):
        self.size += 1
        self.data[self.size] = (0xFFFFFFFF, n)
        self.nodeToHeap[n] = self.size

    def extractMin(self):
        min_elem = self.data[1]
        self.data[1] = self.data[self.size]
        self.size -= 1

        i = 1
        while i < self.size:
            i_min = i
            if 2 * i <= self.size and self.data[2 * i][0] < self.data[i_min][0]:
                i_min = 2 * i

            if 2 * i + 1 <= self.size and self.data[2 * i + 1][0] < self.data[i_min][0]:
                i_min = 2 * i + 1

            if i == i_min:
                break

            self.nodeToHeap[self.data[i_min][1]] = i
            self.nodeToHeap[self.data[i][1]] = i_min
            self.data[i_min], self.data[i] = self.data[i], self.data[i_min]
            i = i_min

        return min_elem

    def updateKey(self, n, key):
        i = self.nodeToHeap[n]
        self.data[i] = (key, self.data[i][1])
        while i > 1 and self.data[i // 2][0] > self.data[i][0]:
            self.nodeToHeap[self.data[i // 2][1]] = i
            self.nodeToHeap[self.data[i][1]] = i // 2
            self.data[i // 2], self.data[i] = self.data[i], self.data[i // 2]
            i //= 2


class HQHeap:
    def __init__(self, nodes):
        self.nodes = nodes
        self.h = []

    def insert(self, n):
        pass

    def updateKey(self, n, k):
        heapq.heappush(self.h, (k, n))

    def extractMin(self):
        return heapq.heappop(self.h)

def main():
    with open('dijkstra.in') as fp:
        nodes, _ = list(map(int, fp.readline().split()))
        adj = {i: [] for i in range(1, nodes + 1)}
        for line in fp:
            u, v, n = list(map(int, line.split()))
            adj[u].append((v, n))

    H = HQHeap(nodes)
    dist = [0, ] * (nodes + 1)

    for i in range(1, nodes + 1):
        H.insert(i)
        dist[i] = 0xFFFFFFFF

    H.updateKey(1, 0)
    dist[1] = 0

    while H.h:
        u = H.extractMin()
        if u[0] != dist[u[1]]:
            continue

        dist[u[1]] = u[0]
        for v in adj[u[1]]:
            if dist[u[1]] + v[1] < dist[v[0]]:
                dist[v[0]] = dist[u[1]] + v[1]
                H.updateKey(v[0], dist[v[0]])

    print(' '.join(list(map(lambda x: str(x) if x != 0xFFFFFFFF else '0',
                            dist[2:nodes+1]))),
          file=open('dijkstra.out', 'w'))

if __name__ == '__main__':
    main()