Estoy creando todas las permutaciones posibles compuestas de tres elementos desde 0 hasta un número dado usando esto:

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

La condición evita tener (0, 1, 2) y (2, 1, 0) "ángulos" al mismo tiempo en mi lista, lo que ya es genial. Ahora, necesito separar esta lista en grupos más pequeños compuestos de "ángulos" que tienen el mismo elemento central. De esta manera tendría:

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

Y así sucesivamente.

¿Podrías ayudarme por favor?

2
Hugo Santos Silva 11 may. 2016 a las 18:46

4 respuestas

La mejor respuesta

Puede considerar 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)]
]
0
schwobaseggl 11 may. 2016 a las 16:26

Si bien las respuestas anteriores proporcionan soluciones excelentes para un problema más general, con solo 3 elementos para elegir una lista, la comprensión es suficiente.

atoms = [0, 1, 2, 3]

angles = [[(a, b, c) for i, a in enumerate(atoms, start=1) if a != b
    for c in atoms[i:] if c != b] for b in atoms]


# [[(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)]]

También es aproximadamente un 40% más rápido que las permutaciones + enfoque grupal en mi máquina.

1
hilberts_drinking_problem 12 may. 2016 a las 13:43

Puede usar defaultdict para agrupar las permutaciones:

from collections import defaultdict

angles = defaultdict(list)

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

Si len(atoms) es 4, obtendría el siguiente resultado:

defaultdict(<type 'list'>, {
    0: [(1, 0, 2), (1, 0, 3), (2, 0, 3)], 
    1: [(0, 1, 2), (0, 1, 3), (2, 1, 3)], 
    2: [(0, 2, 1), (0, 2, 3), (1, 2, 3)], 
    3: [(0, 3, 1), (0, 3, 2), (1, 3, 2)]
})
2
niemmi 11 may. 2016 a las 15:51

La función itertools.groupby se puede usar para crear esas listas que contienen el mismo elemento central, pero primero debe ordenar la lista para que las permutaciones con el mismo elemento central estén una al lado de la otra. Para hacer eso, debe pasar tanto sort como groupby una función clave que mira el elemento central. Una forma de hacerlo es así:

def keyfunc(s):
    return s[1]

O, como una lambda:

keyfunc = lambda s: s[1]

O simplemente puede usar itemgetter del módulo operator, que es breve y significativamente más rápido que usar una función lambda o def.

El siguiente código se basa en su código, pero crea la lista inicial en una comprensión de la lista. Luego coloca los grupos en un diccionario, con el elemento central como la tecla dict.

from itertools import permutations, groupby
from operator import itemgetter

atoms = 'abcd'
perms = permutations(range(len(atoms)), 3)
angles = [list(u) for u in perms if u[0] < u[-1]]

keyfunc = itemgetter(1)

angles.sort(key=keyfunc)
print(angles)

groups = {k: list(g) for k, g in groupby(angles, keyfunc)}
print(groups)

salida

[[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]]
{0: [[1, 0, 2], [1, 0, 3], [2, 0, 3]], 1: [[0, 1, 2], [0, 1, 3], [2, 1, 3]], 2: [[0, 2, 1], [0, 2, 3], [1, 2, 3]], 3: [[0, 3, 1], [0, 3, 2], [1, 3, 2]]}
1
PM 2Ring 12 may. 2016 a las 13:23