Topologically sort DAG into batches

Question

I'm looking for an algorithm to sort a DAG into batches of vertices, such that no vertices in a batch have an edge between them, and the batches are sorted in an order such that no vertex in a batch has an edge to a batch after it in the sorting order. I can't seem to come up with an efficient modification of a generic topological sort, because to discover what batch to put a vertex in, you must fill all batches before the one it belongs in. Effectively:

batches = []
while vertices:
    batch = []
    for v in vertices.copy():
        for v' in v.edges:
            if v' in vertices:
                break
        else:
            batch.append(v)
            vertices.remove(v)
    batches.append(batch)

However, this algorithm is O(n^2) for a reverse sorted 'linear' graph with a hash table for vertex lookup.

Sorry for the pythonic pseudocode.

Dumfound

For graph algorithms, try the library NetworkX. – Osrick

Answer 1

I wrote batching-toposort to solve this for one of my own projects. It is a modification of Kahn's Algorithm with the following performance characteristics:

• Time complexity O(|V| + |E|) for V task vertices and E dependency edges
• Space complexity O(|V|)

The library is in JS, but the core algorithm in pseudocode is as follows; I'm sure you can translate it to Python:

given G = adjacency list of tasks and dependents (~O(1) lookup):

let N = map from tasks to in-degree counts (~O(1) lookup / update)
let L = [] (empty output list) (~O(1) append)
let R1 = list of root tasks (~O(1) addition, ~O(n) iteration)

while R1 is nonempty
    append R1 to L
    let R2 = [] (empty list for next batch) (~O(1) append)
    for each task T in R1
        for each dependent D of T (as per G)
            decrement in-degree count for D (in N)
            if D's in-degree (as per N) is 0
                add D to R2
    R1 = R2

return L

Ring