#!/usr/bin/env python
#coding=utf-8
from heapq import heappop, heappush
def shortest_path(self, start_node, end_node, weight_func=lambda edge: 1):
"""Return the shortest path and its overall weight.
weight_func -- Functions that maps an edge to a weight value, the
default function maps all edges to 1.
"""
## this is basically Dijkstra's algorithm
# store the shortest path to all nodes,
# the value tuples contain the path length and the previous node
shortest_paths = {start_node: (0, None)}
# we use this list as a priority cue with heapq
# the tuples contain the overall path length and the edge itself
edge_heap = []
for edge in start_node.out_edges:
heappush(edge_heap, (weight_func(edge), edge))
while edge_heap:
path_weight, edge = heappop(edge_heap)
if ((edge.head not in shortest_paths) or
(shortest_paths[edge.head][0] > path_weight)):
shortest_paths[edge.head] = (path_weight, edge)
# if we already visited this node then there may be edge
# duplicates in the heap, but this is no problem because
# the newer edges are guaranteed to come first
for out_edge in edge.head.out_edges:
heappush(edge_heap, (path_weight + weight_func(out_edge),
out_edge))
if end_node not in shortest_paths:
err = ("The is no connection from node %s" % str(start_node) +
" to node %s." % str(end_node))
raise NoPathGraphException(err)
# assemble the shortest path from end to start
path_weight = shortest_paths[end_node][0]
path_edges = [shortest_paths[end_node][1]]
current_node = path_edges[-1].tail
while current_node is not start_node:
path_edges.append(shortest_paths[current_node][1])
current_node = path_edges[-1].tail
return path_edges[::-1], path_weight