graph.py

from collections import deque

class Edge:
    def __init__(self, src, tgt):
        self.source = src
        self.target = tgt
        
    def raw(self):
        return (self.source, self.target)
    
    def flip(self):
        return Edge(self.target, self.source)
    
    def __repr__(self):
        return repr(self.raw())
    
    def __hash__(self):
        return hash(self.raw())
    
    def __eq__(self, other):
        return self.raw() == other.raw()

################################################################################
# Directed Adjacency List
################################################################################
    
class DirectedAdjList:
    def __init__(self, edge_list=[], vertex_label=None,
                 vertex_text=None,
                 edge_label=None, edge_color=None):
        self.out = {}
        self.ins = {}
        self.vertex_label = vertex_label
        if vertex_text:
            self.vertex_text = vertex_text
        elif vertex_label:
            self.vertex_text = self.vertex_label
        else:
            self.vertex_text = lambda v: str(v)
        self.edge_label = edge_label
        self.edge_color = edge_color
        self.edge_set = set()
        for e in edge_list:
          if isinstance(e, Edge):
            self.add_edge(e.source, e.target)
          else:
            self.add_edge(e[0], e[1])

    def edges(self):
        return self.edge_set
    
    def vertices(self):
        #return range(0, self.num_vertices())
        return self.out.keys()
    
    def num_vertices(self):
        return len(self.out)
    
    def adjacent(self, u):
        self.add_vertex(u)
        return self.out[u]

    def add_vertex(self, u):
        if u not in self.out:
            self.out[u] = []
            self.ins[u] = []
    
    def add_edge(self, u, v):
        self.add_vertex(u)
        self.add_vertex(v)
        self.out[u].append(v)
        self.ins[v].append(u)
        edge = Edge(u,v)
        self.edge_set.add(edge)
        return edge

    def out_edges(self, u):
        for v in self.out[u]:
            yield Edge(u, v)
            
    def in_edges(self, v):
        for u in self.ins[v]:
            yield Edge(u, v)

    def has_edge(self, u, v):
        return Edge(u,v) in self.edge_set

    def remove_edge(self, u, v):
        self.out[u].remove(v)
        self.ins[v].remove(u)
        self.edge_set.remove(Edge(u,v))
 
    def name(self, u):
        if self.vertex_label:
            return self.vertex_label(u)
        else:
            return str(u)
    
    def named_edge(self, e):
        return (self.name(e.source), self.name(e.target))
        
    def label(self, e):
        if self.edge_label:
            return self.edge_label(e)
        else:
            return ""

    def color(self, e):
        if self.edge_color:
            return self.edge_color(e)
        else:
            return "black"
        
    def show(self, engine='neato'):
      from graphviz import Digraph
      dot = Digraph(engine=engine)
      for u in self.vertices():
          dot.node(self.name(u), self.vertex_text(u))
      for e in self.edges():
          dot.edge(self.name(e.source), self.name(e.target), label=self.label(e),
                   color=self.color(e), len='1.5')
      return dot

class UEdge(Edge):
    def raw(self):
        return set([self.source, self.target])
        
    def __hash__(self):
        return hash(self.source) + hash(self.target)
    
    def __eq__(self, other):
        return self.raw() == other.raw()


################################################################################
# Undirected Adjacency List
################################################################################

class UndirectedAdjList(DirectedAdjList):

    def add_edge(self, u, v):
        self.add_vertex(u)
        self.add_vertex(v)
        self.out[u].append(v)
        self.out[v].append(u)
        edge = UEdge(u,v)
        self.edge_set.add(edge)
        return edge

    def remove_edge(self, u, v):
        self.out[u] = [w for w in self.out[u] if w != v]
        self.out[v] = [w for w in self.out[v] if w != u]
        self.edge_set.remove(UEdge(u,v))
        
    def out_edges(self, u):
        for v in self.out[u]:
            yield UEdge(u, v)

    def in_edges(self, v):
        for u in self.out[v]:
            yield UEdge(u,v)
            
    def has_edge(self, u, v):
        return UEdge(u,v) in self.edge_set

    def remove_edge(self, u, v):
        self.out[u] = [w for w in self.out[u] if w != v]
        self.edge_set.remove(UEdge(u,v))
                
    def show(self):
      from graphviz import Graph
      dot = Graph(engine='neato')
      for u in self.vertices():
        dot.node(self.name(u))
      for e in self.edges():
        dot.edge(self.name(e.source), self.name(e.target), len='1.5')
      return dot


################################################################################
# Topological Sort
################################################################################

def topological_sort(G: DirectedAdjList) -> DirectedAdjList:
    in_degree = {u: 0 for u in G.vertices()}
    for e in G.edges():
        in_degree[e.target] += 1
    queue = deque()
    for u in G.vertices():
        if in_degree[u] == 0:
            queue.append(u)
    topo = []
    while queue:
        u = queue.pop()
        topo.append(u)
        for v in G.adjacent(u):
            in_degree[v] -= 1
            if in_degree[v] == 0:
                queue.append(v)
    return topo  

def transpose(G: DirectedAdjList) -> DirectedAdjList:
    G_t = DirectedAdjList()
    for e in G.edges():
        G_t.add_edge(e.target, e.source)
    return G_t