gistfile1.py

class Graph(object):
    def __init__(self):
        self.nodes = {}
        """
        self.nodes is a mapping between a node and a mapping of nodes that
        can be reached with a set of possible source nodes:
        
        self.nodes = {
            'fromNodeA': {
                'toNodeB': set(['ifComingFromC', 'ifComingFromD', ...]),
                'toNodeE': set(['ifComingFromF']),
            },
            ...
        }
        
        Each node is represented by a unique string.
        """

    def addEdge(self, (node1, comingFrom1), (node2, comingFrom2)):
        """
        Adds an edge to this graph between node1 and node2. This edge can be
        walked in direction node1 -> node2 only if coming from a node in
        comingFrom1, and respectively, in direction node2 -> node1 only if
        coming from a node in comingFrom2.
        
        Complexity: O(1)
        
        Note: comingFrom1 and comingFrom2 can be any iterable, but only
              set/frozenset types allow for the best O complexity in the
              following methods.
        """
        # Add the nodes for both directions, as the graph is not directed.
        self.nodes.setdefault(node1, {})[node2] = comingFrom1
        self.nodes.setdefault(node2, {})[node1] = comingFrom2

    def isEdge(self, comingFrom, passingBy, goingTo):
        """
        Returns True if there is an edge between passingBy and goingTo
        accessible when coming from comingFrom.
        
        Complexity: O(1)
        
        Note: compelxity may be O(n) where n is the number of directions if
              the direction type does not support O(1) lookup.
        """
        try:
            return comingFrom in self.nodes[passingBy][goingTo]
        except KeyError:
            # passingBy not in self.nodes
            # - or -
            # goingTo is not in self.nodes[passingBy]
            return False

    def destinations(self, comingFrom, passingBy):
        """
        Returns a set containing all possible destinations (i.e. edges between
        passingBy and all other nodes in the graph) when coming from
        comingFrom.
        
        Complexity: O(n) where n is the number of edges exiting from passingBy.
        """
        dests = set()
        try:
            for node, directions in self.nodes[passingBy].iteritems():
                if comingFrom in directions:
                    dests.add(node)
        except KeyError:
            pass
        
        return dests

    def sources(self, passingBy, goingTo):
        """
        Returns a set containing all possible source nodes which allow to
        walk the edge between passingBy and goingTo.
        
        Complexity: O(n) where n is the number of edges entering passingBy.
        """
        return self.destinations(goingTo, passingBy)