4_5.py

#!/usr/bin/python

from Queue import Queue as Q

class UndirectedGraph:
    def __init__(self, vertices, edges):
        self.vertices = vertices
        self.edges = edges
        if not self.validate_edges(self.edges):
            print "Invalid Graph"
            exit(1)

    def validate_edges(self, edges):
        for (v1, v2) in edges:
            if (v1 not in self.vertices or
                v2 not in self.vertices):
                return False

        return True

    def edges_for_vertex(self, vertex):
        return map(lambda (v1,v2): v2 if vertex == v1 else v1,
                   filter(lambda (v1,v2): vertex == v1 or vertex == v2,
                          self.edges))

    def find_paths(self, start, end):
        dist = {}
        count = {}
        frontier = Q()
        frontier.put(start)
        dist[start] = 0
        count[start] = 1
        while not frontier.empty():
            u = frontier.get()
            if u == end:
                print "Found {} shortest paths from {} to {}".format(
                    count[u], start, end
                )
            print "Processing {}".format(u)
            for v in self.edges_for_vertex(u):
                try:
                    # Already visited
                    if dist[v] == dist[u]+1:
                        count[v] = count[v] + count[u]
                    elif dist[v] > dist[u]+1:
                        dist[v] = dist[u] + 1
                        count[v] = count[u]
                except KeyError:
                    print "Adding {}".format(v)
                    frontier.put(v)
                    dist[v] = dist[u] + 1
                    count[v] = count[u]
        print dist, count

def main():
    g1 = UndirectedGraph(['A', 'B', 'C', 'D', 'E', "F"],
                         [('A', 'B'), ("D", "E"),
                          ("A", "D"), ("C", "E"),
                          ("B", "D"), ("B", "E"),
                          ("A", "F"), ("F", "E")])
    g1.find_paths('A', 'D') # Should be length: 1, count: 1
    g1.find_paths('A', 'C') # Should be length: 3, count: 2
    g1.find_paths('A', 'A') # Should be length: 3, count: 2


if __name__ == "__main__":
    main()