shortest_path.py

import heapq


class Vertex:
    def __init__(self, val):
        self.val = val
        self.adjs = {}

    def add_edge(self, j, val):
        self.adjs[j] = val


class Graph:
    def __init__(self, v):
        self.v = v
        self.vertexes = []
        for i in range(v):
            self.vertexes.append(Vertex(i))

    def add_edge(self, i, j, val):
        self.vertexes[i].add_edge(j, val)

    def print_path(self, i, prev):
        if i:
            self.print_path(prev[i], prev)
            print i

    def shortest_path(self, origin, to):
        visited = [False] * self.v
        distance = [None] * self.v
        prev = [None] * self.v
        q = []
        distance[origin] = 0
        heapq.heappush(q, (0, origin))
        while q:
            item = heapq.heappop(q)
            cost, i = item[0], item[1]
            if visited[i]:
                continue
            if i == to:

                return cost
            visited[i] = True
            for adj, val in self.vertexes[item[1]].adjs.items():
                if distance[adj] is None or cost+val < distance[adj]:
                    prev[adj] = i
                    total_cost = cost+val
                    distance[adj] = total_cost
                    heapq.heappush(q, (total_cost, adj))

        return None

g = Graph(4)
g.add_edge(0, 1, 20)
g.add_edge(0, 2, 10)
g.add_edge(0, 3, 55)
g.add_edge(1, 3, 30)
g.add_edge(2, 3, 30)
print g.shortest_path(0, 2)