Dijkstra.py

import collections
import heapq

Connection = collections.namedtuple('Connection', 'node weight')
Route = collections.namedtuple('Route', 'cost path')

class Heap(object):
    def __init__(self):
        self._values = []

    def push(self, value):
        heapq.heappush(self._values, value)

    def pop(self):
        return heapq.heappop(self._values)

    def __len__(self):
        return len(self._values)

class Graph:
    def __init__(self):
        self.__nodes = collections.defaultdict(set)

    def add_connection(self, origin, destiny, weight):
        self.__nodes[origin].add(Connection(node=destiny, weight=weight))
        self.__nodes[destiny].add(Connection(node=origin, weight=weight))

    def get_connections(self, node):
        return self.__nodes[node]

    def dijkstra(self, origin, destiny):
        h = Heap()
        h.push(Route(cost=0, path=[origin]))
        visited = set()

        while h:
            price, path = h.pop()
            n_node = path[-1]

            if n_node in visited:
                continue

            if n_node == destiny:
                return price, path
            
            for n in self.get_connections(n_node):
                if n[0] not in visited:
                    r = Route(cost=price+n[1], path=path + [n[0]])
                    h.push(r)

            visited.add(n_node)

        raise Exception("unexpected")