Search code examples
python-3.xlist-comprehensionpython-itertools

nested comprehensions with itertools


I want to create a list with the vertices of all paths of length 2 from a set of 5 vertices.
What needs to be done is clear: list all permutations of every subset of 3 vertices.

itertools provides methods for generating all subsets via itertools.combinations and for generating all permutations.

Generating the comprehensions separately is easy:

import itertools as iter  
p = [i for i in iter.permutations(range(5))]  
q = [j for j in iter.combinations(range(5),3)]  

but how can I achieve something like

import itertools as iter
c= [ j for sub in list(list(iter.permutations(i)) for i in list(iter.combinations(range(5),3))) for j in sub]

in a simpler, i.e. more compact, way?


Solution

  • As @JeffH noted in a comment, passing a size argument to permutations appears to directly do what you asked for:

    import itertools
    vs = ['v1', 'v2', 'v3', 'v4']
    L1 = list(itertools.permutations(vs, 3))
    print(L1)
    

    which displays

    [('v1', 'v2', 'v3'), ('v1', 'v2', 'v4'), ('v1', 'v3', 'v2'), ('v1', 'v3', 'v4'),
     ('v1', 'v4', 'v2'), ('v1', 'v4', 'v3'), ('v2', 'v1', 'v3'), ('v2', 'v1', 'v4'),
     ('v2', 'v3', 'v1'), ('v2', 'v3', 'v4'), ('v2', 'v4', 'v1'), ('v2', 'v4', 'v3'),
     ('v3', 'v1', 'v2'), ('v3', 'v1', 'v4'), ('v3', 'v2', 'v1'), ('v3', 'v2', 'v4'),
     ('v3', 'v4', 'v1'), ('v3', 'v4', 'v2'), ('v4', 'v1', 'v2'), ('v4', 'v1', 'v3'),
     ('v4', 'v2', 'v1'), ('v4', 'v2', 'v3'), ('v4', 'v3', 'v1'), ('v4', 'v3', 'v2')]
    

    But if permutations() didn't have that second argument, you could get the same effect (although perhaps in a different order) via other 1-statement compositions; e.g.,

    L2 = list(itertools.chain.from_iterable(
              map(itertools.permutations,
                  itertools.combinations(vs, 3))))
    assert len(L1) == len(L2)
    assert set(L1) == set(L2)