graph-contract.py

def contract(g, f):
    """
    Returns a directed acyclic graph with all vertices from the input DAG for
    which f() returns True.
    """
    vertices = [v for v in g.keys() if f(v)]
    result = defaultdict(set, ((v, set()) for v in vertices))

    stack = zip(vertices, vertices)
    while len(stack):
        v, base = stack.pop()
        for w in g.get(v, set()):
            if f(w):
                result[base].add(w)
            else:
                stack.insert(0, (w, base))

    return result