Instantly share code, notes, and snippets.

# GaretJax/gist:1030049 Created Jun 16, 2011

 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)