Search code examples
pythonlistcombinationspython-itertools

Python combinations of multiple list of different sizes


Am trying to swap items between multiple lists and I wanted to know if there is any method to generate combinations between multiple list of different size?

For example, I have this 3 lists:

a = [(0, 0), (1, 0), (2, 0)]
b = [(0, 2), (1, 2), (2, 2)]
c = [(0, 3), (1, 3)]

Expected result:

a : [(0, 3), (0, 2), (0, 0)]
b : [(1, 3), (1, 2), (1, 0)]
c : [(2, 2), (2, 0)]

a : [(0, 3), (0, 2), (0, 0)]
b : [(1, 3), (1, 2), (2, 0)]
c : [(2, 2), (1, 0)]

...

a : [(0, 3), (0, 2)]
b : [(1, 3), (1, 2), (0, 0)]
c : [(2, 2), (2, 0), (1, 0)]

I found this code here (python combinations of multiple list):

import itertools as it
import numpy as np

a = [(0, 0), (1, 0), (2, 0)]
b = [(0, 2), (1, 2), (2, 2)]
c = [(0, 3), (1, 3)]

def combination(first, *rest):
    for i in it.product([first], *(it.permutations(j) for j in rest)):
        yield tuple(zip(*i))

for i in combination(c, b, a):
    print("a :", list(i[0]))
    print("b :", list(i[1]))
    print("c :", list(i[2]))

It works perfectly fine if the list are the same size.


Solution

  • Try

    1. adding None to your lists so that they all have the same length,
    2. use sympy.utilities.iterables.multiset_permutations instead of, it.permutations, and
    3. finally filter out None values from the output.

    That should generalize in a natural way your approach for lists of equal sizes:

    import itertools as it
    from sympy.utilities.iterables import multiset_permutations
    
    a = [(0, 0), (1, 0), (2, 0)]
    b = [(0, 2), (1, 2), (2, 2)]
    c = [(0, 3), (1, 3), None]
    
    def combination(first, *rest):
        for i in it.product([first], *(multiset_permutations(j) for j in rest)):
            yield tuple(zip(*i))
    
    for i in combination(c, b, a):
        print("a :", [val for val in i[0] if val])
        print("b :", [val for val in i[1] if val])
        print("c :", [val for val in i[2] if val])