graph.py

#!/usr/bin/env python


class Node(object):
    """A node in a directed graph."""
    def __init__(self, name, outs):
        """Initializes a Node.

        Args:
            name: The name of this node.
            outs: Nodes with an edge leading out from this node.
        """
        self.name = name
        self.outs = outs

    def __repr__(self):
        return self.name


class Graph(object):
    """A directed graph."""
    def __init__(self, nodes):
        """Initializes a Graph.

        Args:
            nodes: A list of nodes in this graph.
        """
        self.nodes = nodes

    def find_cycle(self):
        known_acyclic = set()
        for node in self.nodes:
            cycle = self.find_cycle_from_node(node, [], known_acyclic)
            if cycle:
                return cycle
        return None

    def find_cycle_from_node(self, node, path, known_acyclic):
        """Finds a cycle from a given node if there is one.

        Recursively searches every path through the graph, looking for a cycle.

        Args:
            node: Current node to check for cycles.
            path: A list containing the acyclic path to the current node from the
                first node.
            known_acyclic: Once we know that no path from a particular node has
                a cycle, we can ignore any paths going through that node.

        Returns:
            A list of nodes that make up a cycle or None if no cycle exists.
            The list will begin and end with the same node, i.e. [A, B, A].
        """

        if node in known_acyclic:
            return None

        path.append(node)
        if node in path[:-1]:
            # If path is [A, B, Y, Z, Y], return [Y, Z, Y] (ignore the A, B steps
            # leading to the cycle).
            return path[path.index(node):]

        for out_node in node.outs:
            cycle = self.find_cycle_from_node(out_node, path, known_acyclic)
            if cycle:
                return cycle

        path.pop()
        known_acyclic.add(node)
        return None


### UNIT TEST

import unittest

def _cycle_test(*paths):
    """Test find_cycle on a test graph, indicated by the list of paths.

    Args:
        paths: A list of paths. Each path is a list of node names. For each
        successive two nodes in a path (X, Y), an edge exists from X to Y.
    """
    node_names = sorted({n for p in paths for n in p})
    nodes = {n: Node(n, []) for n in node_names}
    for p in paths:
        for i in range(len(p) - 1):
            nodes[p[i]].outs.append(nodes[p[i + 1]])
    graph = Graph([nodes[n] for n in node_names])
    cycle = graph.find_cycle()
    if cycle is None:
        return None
    return [n.name for n in cycle]

class GraphCycleTest(unittest.TestCase):

    def test_cycle_detection(self):
        (A, B, C, D) = "ABCD"
        self.assertEqual(_cycle_test([A, B, A]), [A, B, A])
        self.assertEqual(_cycle_test([A, B, C, A]), [A, B, C, A])
        self.assertEqual(_cycle_test([A, C], [B, C], [D]), None)
        self.assertEqual(_cycle_test([A, B], [B, C, B], [C, D]), [B, C, B])