Python programs / Graphs / Topological sort

Python Program for Topological Sort (Kahn's Algorithm)

Advanced Graphs Time O(V + E) Space O(V)

Orders courses so that each comes after the courses it needs, with Kahn's algorithm: count each one's prerequisites, then use a queue.

A topological sort puts the nodes of a directed graph in a line so that every edge points forward: if A has to come before B, A is earlier in the line. It answers questions like "in what order can I take these courses?" or "in what order do these build steps have to run?". In the program, Data needs both Python and Maths, so it can only come after both. Kahn's algorithm builds the order by taking, again and again, a node that nothing still waits for. An order exists only when the graph has no cycle: if A needs B and B needs A, neither can go first.

Program

from collections import deque


def topological_sort(graph):
    """The nodes in an order where every edge points forward."""
    in_degree = {node: 0 for node in graph}
    for node in graph:
        for after in graph[node]:
            in_degree[after] += 1
    queue = deque(node for node in graph if in_degree[node] == 0)
    order = []
    while queue:
        node = queue.popleft()
        order.append(node)
        for after in graph[node]:
            in_degree[after] -= 1
            if in_degree[after] == 0:  # all it needs is done
                queue.append(after)
    if len(order) < len(graph):
        return None  # a cycle: some nodes wait on each other
    return order


# Each course, and the courses that build on it
courses = {
    "Python": ["Data", "Web"],
    "Maths": ["Data"],
    "Data": ["ML"],
    "Web": ["Project"],
    "ML": ["Project"],
    "Project": [],
}
print("Course order:", topological_sort(courses))

Open in the editor Download topological_sort.py

Output

Course order: ['Python', 'Maths', 'Web', 'Data', 'ML', 'Project']

How it works

  1. in_degree counts the edges coming into each node: how many courses it needs first. Data needs Python and Maths, so its count is 2.
  2. The nodes with a count of 0 need nothing, so they can go first: they start the queue, a deque from the collections module.
  3. order starts empty. While the queue has nodes, the one at the front goes next in the order.
  4. Now that it's done, every node that needed it needs one fewer. A node whose count reaches 0 has everything it needs, so it joins the queue.
  5. If some nodes never got to 0, they wait on each other in a cycle and no order exists, so the function returns None. Otherwise it returns the order.
  6. The graph is a dictionary: each course is a key, and its list holds the courses that build on it. Python leads to Data and Web; Web and ML both lead to Project.
  7. The program prints an order to take the courses in.

The counts start at Python 0, Maths 0, Data 2, Web 1, ML 1 and Project 2, so the queue starts with Python and Maths:

Step Takes Counts that drop Queue after
1 Python Data 2 → 1, Web 1 → 0 Maths, Web
2 Maths Data 1 → 0 Web, Data
3 Web Project 2 → 1 Data
4 Data ML 1 → 0 ML
5 ML Project 1 → 0 Project
6 Project none empty

Other orders are right too, such as Maths, Python, Data, Web, ML, Project. Which one Kahn's algorithm gives depends on the order of the dict and its lists.

Time and space complexity

TimeO(V + E)
Extra spaceO(V)

Counting the in-degrees reads every node's list once: O(V + E) for V nodes and E edges. Then each node goes into the queue once and comes out once, and when it comes out, each of its edges lowers one count: O(V + E) again. in_degree, the queue and order each hold at most V nodes, so the extra space is O(V).

Other ways to write it

A depth-first search adds a node to order only after every course that builds on it has been added, so the list comes out back to front, and reversing it gives an order. Unlike Kahn's algorithm, this version doesn't notice a cycle.

def topological_sort(graph):
    order = []
    seen = set()

    def visit(node):
        seen.add(node)
        for after in graph[node]:
            if after not in seen:
                visit(after)
        order.append(node)  # after everything that builds on it

    for node in graph:
        if node not in seen:
            visit(node)
    return order[::-1]

courses = {"Python": ["Data", "Web"], "Maths": ["Data"],
           "Data": ["ML"], "Web": ["Project"], "ML": ["Project"],
           "Project": []}
print(topological_sort(courses))

With graphlib

Since Python 3.9 the standard library has a topological sort, graphlib.TopologicalSorter. It takes each node with the nodes it needs first, the other way round from the program's dict, and raises CycleError for a cycle.

from graphlib import TopologicalSorter

needs = {"Data": ["Python", "Maths"], "Web": ["Python"],
         "ML": ["Data"], "Project": ["Web", "ML"]}
print(list(TopologicalSorter(needs).static_order()))

Try it yourself

  • Make Project something Python needs, "Project": ["Python"]. Now the courses wait on each other in a circle, and the function returns None.
  • Move "Maths" above "Python" in the dict. The order becomes Maths, Python, Data, Web, ML, Project: different, and still right.
  • Add a course "Job": [] and put it in Project's list. It comes last.
  • Group the courses into terms: in each round, take every course in the queue at once, as one term. The terms are [['Python', 'Maths'], ['Web', 'Data'], ['ML'], ['Project']].

More graphs programs

Share this program