[Feature-request] Topological sort with keying function that exposes internal state

[kaushikcfd]

Consider the DAG below:

from pytools.graph import compute_topological_order

dag = {"A": {"C"},
       "B": {"E"},
       "C": {"D"},
       "D": set(),
       "E": set(),
       }

colors = {"A": "red",
          "B": "red",
          "C": "blue",
          "D": "red",
          "E": "blue"
          }

sorted_nodes = compute_topological_order(dag, key=lambda x: x)
print(sorted_nodes)                            # ['A', 'B', 'C', 'D', 'E']
print([colors[node] for node in sorted_nodes]) # ['red', 'red', 'blue', 'red', 'blue']

My application needs a valid topological sort that minimizes the jumps from 'red'<->'blue' nodes in the final schedule. And so for this problem, the topological sort order I'm after would be ['A', 'B', 'C', 'E', 'D']. I don't think that's possible with the current interface of compute_topological_order.

[inducer]

Is some greedy thing enough, or do you need an actual minimum?

[kaushikcfd]

A greedy thing seems enough.