Search code examples
pythoncombinationspython-itertools

Generate list combinations in Python - with rules


I have two list:

list1 = ["A", "B", "C", "D", "E"]
list2 = ["AA", "BB", "CC", "DD", "EE"]

I want to create every possible combination, as the following:

    A-B-C-D-E
    A-BB-C-D-E
    A-B-C-DD-E...etc.

The rule is that two similar letter (like A-AA, B-BB) can't be in the combination at the same time, and the order is reversible (A-B-C-D-E and B-A-E-C-D is the same in content so i don't need them both)

How can I manage it with itertools?


Solution

  • TL;DR: map('-'.join, product(*zip(list1, list2))).


    It's relatively simple to do using itertools, but you have to think through each step carefully.

    First, you can zip your two lists together to get a sequence of tuples. Every element in your final result will include exactly one choice from each tuple.

    >>> list1 = ["A", "B", "C", "D", "E"]
    >>> list2 = ["AA", "BB", "CC", "DD", "EE"]
    >>> list(zip(list1, list2))
    [('A', 'AA'), ('B', 'BB'), ('C', 'CC'), ('D', 'DD'), ('E', 'EE')]
    

    Next, we'll want the Cartesian product of each of the five tuples. This gives us the 32 different ways of choosing one of A or AA, then one of B or BB, etc. To do that, we use * to unpack the result of zip into five separate arguments for product.

    >>> from itertools import product
    >>> for x in product(*zip(list1, list2)):
    ...   print(x)
    ...
    ('A', 'B', 'C', 'D', 'E')
    ('A', 'B', 'C', 'D', 'EE')
    ('A', 'B', 'C', 'DD', 'E')
    ('A', 'B', 'C', 'DD', 'EE')
    ('A', 'B', 'CC', 'D', 'E')
    ('A', 'B', 'CC', 'D', 'EE')
    ('A', 'B', 'CC', 'DD', 'E')
    # etc
    

    Once you have the product, each element of the product is a valid argument to '-'.join to create one of the strings in your desired set:

    >>> for x in map('-'.join, product(*zip(list1, list2))):
    ...   print(x)
    ...
    A-B-C-D-E
    A-B-C-D-EE
    A-B-C-DD-E
    A-B-C-DD-EE
    A-B-CC-D-E
    A-B-CC-D-EE
    A-B-CC-DD-E
    A-B-CC-DD-EE
    A-BB-C-D-E
    A-BB-C-D-EE
    A-BB-C-DD-E
    A-BB-C-DD-EE
    A-BB-CC-D-E
    A-BB-CC-D-EE
    A-BB-CC-DD-E
    A-BB-CC-DD-EE
    AA-B-C-D-E
    AA-B-C-D-EE
    AA-B-C-DD-E
    AA-B-C-DD-EE
    AA-B-CC-D-E
    AA-B-CC-D-EE
    AA-B-CC-DD-E
    AA-B-CC-DD-EE
    AA-BB-C-D-E
    AA-BB-C-D-EE
    AA-BB-C-DD-E
    AA-BB-C-DD-EE
    AA-BB-CC-D-E
    AA-BB-CC-D-EE
    AA-BB-CC-DD-E
    AA-BB-CC-DD-EE