Hugo Santos Silva Hugo Santos Silva - 6 months ago 28
Python Question

assembling lists with common central element

I'm creating all the possible permutations composed of three elements from 0 to a given number using this:

for i in itertools.permutations(range(len(atoms)), 3):
if i[0] < i[-1]:
angles = list(i)


The condition avoids having (0, 1, 2) and (2, 1, 0) "angles" at the same time on my list, what is already great. Now, I need to separate this list into smaller groups being composed of "angles" which have the same central element.
In this way I'd have:

A = ([0, 1, 2], [0, 1, 3], [3, 1, 4])...
B = ([0, 2, 3], [0, 2, 4], [3, 2, 4])...


and so on.

Could you please help me out?

Answer

You might consider itertools.groupby:

from itertools import groupby, permutations

perms = filter(lambda x: x[0] < x[-1], permutations(range(4), 3))
key = lambda x: x[1]  # sort and group by second element
angles = [list(g) for k, g in groupby(sorted(perms, key=key), key=key)]
# here, any comprehension can be used, e.g.
# angles = {k: list(g) for k, g in groupby(sorted(perms, key=key), key=key)}
# will produce the dict from @niemmi's answer

>>> angles
[
    [(1, 0, 2), (1, 0, 3), (2, 0, 3)], 
    [(0, 1, 2), (0, 1, 3), (2, 1, 3)], 
    [(0, 2, 1), (0, 2, 3), (1, 2, 3)], 
    [(0, 3, 1), (0, 3, 2), (1, 3, 2)]
]