Random graph and reachable items

[Gribouillis]

I've been playing with generators to write a function that yields the reachable elements in a graph. It is an example of filling a container while iterating on it. It is not particularly difficult but I think some of you might find this interesting

# somegraphtools.py
# Author: Gribouillis for www.python-forum.io
# License: MIT (https://opensource.org/licenses/MIT)
#
# Recent changes:
#   * Added tclosure() function
#   * Added reachable() function

from collections import defaultdict
import functools
from more_itertools import unique_everseen
import random

__version__ = '2019.06.07.2'

def injective(generator):
    """Decorator that transform a generator into an injective generator
    
    An injective generator generates a sequence of items that appear
    only once in the sequence.
    """
    @functools.wraps(generator)
    def wrapper(*args, **kwargs):
        return unique_everseen(generator(*args, **kwargs))
    return wrapper

@injective
def consume_set(to_do):
    """Yield unique items popped from a set until the set is empty.
    
    The set can be updated during iteration.
    """
    try:
        while True:
            yield to_do.pop()
    except KeyError:
        pass

def reachable(start, neighbors, include_start=True):
    """Generates reachable elements in a directed graph.
    
    Arguments:
        * start     : an arbitrary initial element (any python object)
        * neighbors : a function such that neighbors(item) is a sequence
                    of the item's neighbors in the graph.
        * include_start=True : forces the inclusion of start in the
                    result. If this is True, 'start' will be the first
                    item in the result. If this is False, 'start' may
                    be present in the list iff there is a cycle
                    containing start in the graph.
                    
    Returns:
        A injective sequence of all reachable elements from
        element 'start'.
    """
    to_do = set([start] if include_start else neighbors(start))
    for x in consume_set(to_do):
        yield x
        to_do.update(neighbors(x))

def tclosure(edges):
    """Computes the transitive closure of a directed graph.
    
    Arguments:
        * edges : an iterable of pairs (src, dst) representing
            the edges of the graph. All vertices src and dst need
            to be hashable python objects
            
    Return:
        A dictionary  {src : {set of reachable dst}} representing
        the transitive closure of this graph.
    """
    to_do = set(edges)
    left = defaultdict(set)
    right = defaultdict(set)
    for a, b in consume_set(to_do):
        to_do.update((a, c) for c in right[b])
        to_do.update((c, b) for c in left[a])
        right[a].add(b)
        left[b].add(a)
    return right


def random_graph(size, density):
    """Creates a random graph which vertices are integers in range(size).
    
    Arguments:
        * size    : a positive integer.
        * density : a real number in the interval [0, 1]. The expected
                    number of neighbors of a node is density * size.
                    
    Returns:
        A list of length size which item at index k is a set
        of indices representing the neighbors of k in a directed graph.
    """
    return [ set([j for j in range(size) if random.random() < density])
            for i in range(size)]

def main():
    G = random_graph(10, 0.15)
    print(G)
    print(list(reachable(0, G.__getitem__, False)))
    C = tclosure((i, j) for i, s in enumerate(G) for j in s)
    print([C[i] for i in range(len(G))])
    # check correctness of tclosure
    for i in range(len(G)):
        assert set(reachable(i, G.__getitem__, False)) == C[i]

if __name__ == '__main__':
    main()

Example output

Output:
[{4, 7}, {1, 3}, {8, 1, 2, 7}, set(), {6}, {1}, set(), {6}, {8, 1, 2}, set()]
[0, 4, 6, 7]

[Gribouillis]

In version 2019.06.07.2 above, I added an algorithm to compute the transitive closure of the directed graph by reusing the same "set consuming" function.