Cod sursa(job #2307047)

Utilizator 24601Dan Ban 24601 Data 23 decembrie 2018 17:03:46
Problema Algoritmul lui Dijkstra Scor 20
Compilator py Status done
Runda Arhiva educationala Marime 1.99 kb
class Heap:
    def __init__(self, nodes):
        self.data = [(0, 0), ] * (nodes + 1)
        self.nodeToHeap = {}
        self.size = 0

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

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

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

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

        if i == i_min:
            break
        else:
            H.nodeToHeap[H.data[i_min][1]] = i
            H.nodeToHeap[H.data[i][1]] = i_min
            H.data[i_min], H.data[i] = H.data[i], H.data[i_min]
            i = i_min

    return min_elem

def updateKey(H, n, key):
    i = H.nodeToHeap[n]
    H.data[i] = (key, H.data[i][1])
    while i > 1:
        if H.data[i // 2][0] > H.data[i][0]:
            H.nodeToHeap[H.data[i // 2][1]] = i
            H.nodeToHeap[H.data[i][1]] = i // 2
            H.data[i // 2], H.data[i] = H.data[i], H.data[i // 2]
            i //= 2
        else:
            break

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 = Heap(nodes)
    dist = [0, ] * (nodes + 1)

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

    updateKey(H, 1, 0)

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

    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()