Depth-First Search

There are a couple of ways of implementing topological sort using a depth-first search. One approach is recursive and another uses a stack. I will cover the stack-based approach, but feel free to try the other implementation and compare.

With depth-first search we first add our "start nodes", the ones with no outgoing edges, to the stack. We also use a set or dictionary to keep track of which nodes we've already seen in the search. We then proceed as follows. We pop a node from the stack and check if we've visited it before. If we haven't, then we add its parents to the stack (the nodes it has incoming edges from) and add the popped node to the rightmost place in the topological sort. If we have seen the node though, then this means we sorted it wrong in the first place and need to remove it from its original place from the ordering. In Python, we can delete an element from a list using the del keyword. Confused yet? I don't blame you. Let's dive into an exercise to solidify our understanding.

Exercise: Download the attached file and implement your own topological sort. How does it compare to the reference solution below? Can you make it faster? The exercise file includes code for testing your solution's speed against that of the one written by the Autograd authors.

View Solution

def topological_sort(end_node):
    stack = [end_node]
    visited = {}
    topo_sorted = []
    while stack:
        node = stack.pop()
        if node in visited:
            del topo_sorted[topo_sorted.index(node)]
        visited[node] = True
        parents = node.parents
        stack.extend(parents)
        topo_sorted.append(node)
    return topo_sorted