CODEWITHSUNDEEP
python-3.xlist-comprehensionpython-itertools
Manfred Weis
Manfred Weis

Reputation: 836

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?

Upvotes: 0

Views: 61

Answers (1)

Tim Peters
Tim Peters

Reputation: 70602

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)

Upvotes: 2

Related Questions