one-file-projects/graph.py

#!/usr/bin/env python

class Infinity:
    def __eq__(self,other):
        return isinstance(other,Infinity)
    def __gt__(self,other):
        return not isinstance(other, Infinity)
    def __lt__(self,other):
        return False

    def __add__(self,other):
        return self
    def __radd__(self,other):
        return self

    def __repr__(self):
        return "inf"

class Graph:
    def __init__(self):
        self.nodes = set()
        self.edges = set()

    def add_node(self,node):
        self.nodes.add(node)

    def add_edge(self,nodefrom, nodeto, extra=None, both=False):
        self.require_node(nodefrom)
        self.require_node(nodeto)

        self.edges.add( (nodefrom,nodeto,extra) )

        if both:
            self.add_edge(nodeto, nodefrom, extra=extra, both=False)

    def all_edges_of(self,nodefrom):
        return filter(lambda e: e[0] == nodefrom, self.edges)

    def has_node(self, node):
        return node in self.nodes

    def require_node(self,node):
        if not self.has_node(node):
            raise Exception("Node %s does not exist" % repr(node))


def dijkstra(graph,startnode):
    graph.require_node(startnode)
    d = { node: Infinity() for node in graph.nodes }
    d[startnode] = 0
    V = [startnode]

    while len(V) > 0:
        node = min(V, key=lambda e: d[e])
        V.remove(node)
        for _, neighbor, distance in graph.all_edges_of(node):
            if d[neighbor] > d[node] + distance:
                d[neighbor] = d[node] + distance
                V.append(neighbor)

    return d

def bfs(graph, startnode):
    graph.require_node(startnode)
    visited = set()
    not_visited = [startnode]
    while len(not_visited) > 0:
        node = not_visited.pop(0)
        yield node
        visited.add(node)
        for _, neighbor, extra in graph.all_edges_of(node):
            if neighbor not in visited and neighbor not in not_visited:
                not_visited.append(neighbor)

def dfs(graph, node, visited=None):
    graph.require_node(node)
    if visited is None:
        visited = []
    yield node
    visited.append(node)
    for _, neighbor, extra in graph.all_edges_of(node):
        if neighbor not in visited:
            for node in dfs(graph, neighbor, visited=visited):
                yield node

def main():
    g = Graph()
    for i in range(1,6):
        g.add_node(i)

    g.add_edge(1,2,10)
    g.add_edge(1,4,30)
    g.add_edge(1,5,100)
    g.add_edge(2,3,50)
    g.add_edge(3,5,10)
    g.add_edge(4,3,20)
    g.add_edge(4,5,60)

    d = dijkstra(g,1)

    print(d)

    print( list(bfs(g,1)) )
    print( list(dfs(g,1)) )

    g2 = Graph()
    for i in range(1,10):
        g2.add_node(i)

    for i in range(1,5):
        g2.add_edge(i,2*i)
        g2.add_edge(i,2*i+1)

    print( list(bfs(g2,1)) )
    print( list(dfs(g2,1)) )



if __name__ == '__main__':
    main()