CODEWITHSUNDEEP
pythoncombinationspython-itertoolsnested-lists
ABarrier
ABarrier

Reputation: 962

Python nested lists, all combinations

I have nested lists of the form:

['A', 'B', ['a0', 'a1', 'a2'], 'C', 'D', ['a0', 'a1', 'a2'], 'E', 'F', ['a0', 'a1', 'a2']]

and I would like to generate all the combinations with respect to the sublist ['a0', 'a1', 'a2'] of the form:

['A', 'B', ['a0'], 'C', 'D', ['a0'], 'E', 'F', ['a0']]
['A', 'B', ['a0'], 'C', 'D', ['a0'], 'E', 'F', ['a1']]
['A', 'B', ['a0'], 'C', 'D', ['a0'], 'E', 'F', ['a2']]
['A', 'B', ['a0'], 'C', 'D', ['a1'], 'E', 'F', ['a0']]
['A', 'B', ['a0'], 'C', 'D', ['a1'], 'E', 'F', ['a1']]
['A', 'B', ['a0'], 'C', 'D', ['a1'], 'E', 'F', ['a2']]
['A', 'B', ['a0'], 'C', 'D', ['a2'], 'E', 'F', ['a0']]
['A', 'B', ['a0'], 'C', 'D', ['a2'], 'E', 'F', ['a1']]
['A', 'B', ['a0'], 'C', 'D', ['a2'], 'E', 'F', ['a2']]
.             .                 .
.             .                 .
.             .                 .
['A', 'B', ['a2'], 'C', 'D', ['a2'], 'E', 'F', ['a2']]

etc. In total 27 27 lists. I know I have to use itertools package but I cannot figure out how. Any idea welcomed.

Upvotes: 1

Views: 593

Answers (2)

Dani Mesejo
Dani Mesejo

Reputation: 61930

You could do, (if there is only one single nested level):

from itertools import product

lst = ['A', 'B', ['a0', 'a1', 'a2'], 'C', 'D', ['a0', 'a1', 'a2'], 'E', 'F', ['a0', 'a1', 'a2']]


def nested_product(ls):
    lst_positions = [l for l in ls if isinstance(l, list)]
    for p in product(*lst_positions):
        it = iter(p)
        yield [e if not isinstance(e, list) else [next(it)] for e in lst]


for pr in nested_product(lst):
    print(pr)

Output (partial)

['A', 'B', ['a0'], 'C', 'D', ['a0'], 'E', 'F', ['a0']]
['A', 'B', ['a0'], 'C', 'D', ['a0'], 'E', 'F', ['a1']]
['A', 'B', ['a0'], 'C', 'D', ['a0'], 'E', 'F', ['a2']]
['A', 'B', ['a0'], 'C', 'D', ['a1'], 'E', 'F', ['a0']]
['A', 'B', ['a0'], 'C', 'D', ['a1'], 'E', 'F', ['a1']]
['A', 'B', ['a0'], 'C', 'D', ['a1'], 'E', 'F', ['a2']]
['A', 'B', ['a0'], 'C', 'D', ['a2'], 'E', 'F', ['a0']]
['A', 'B', ['a0'], 'C', 'D', ['a2'], 'E', 'F', ['a1']]
['A', 'B', ['a0'], 'C', 'D', ['a2'], 'E', 'F', ['a2']]
['A', 'B', ['a1'], 'C', 'D', ['a0'], 'E', 'F', ['a0']]
...

An alternative solution, using yield from (Python 3.3+), is the following:

def build_nested_list(ls, it):
    return [e if not isinstance(e, list) else [next(it)] for e in ls]


def nested_product(ls):
    lst_positions = (li for li in ls if isinstance(li, list))
    yield from (build_nested_list(ls, iter(p)) for p in product(*lst_positions))

Upvotes: 4

Aplet123
Aplet123

Reputation: 35560

If there's only one layer of nesting:

import itertools

def combos(lst):
    # prep a copy without sublists so a little less can be copied
    copy = []
    lindices = []
    for i, v in enumerate(lst):
        if isinstance(v, list):
            # if it's a list we note the index
            lindices.append(i)
            # we append None to the copy since it'll get overwritten anyway
            copy.append(None)
        else:
            copy.append(v)
    ret = []
    # we cartesian product all the list arguments
    # the * is argument destructuring
    for items in itertools.product(*[lst[i] for i in lindices]):
        # we make a copy of the list
        curcopy = copy.copy()
        for i, item in zip(lindices, items):
            # we assign the elements
            curcopy[i] = item
        # we append the copy to the return
        ret.append(curcopy)
    return ret

test = ['A', 'B', ['a0', 'a1', 'a2'], 'C', 'D', ['b0', 'b1', 'b2'], 'E', 'F', ['c0', 'c1', 'c2']]
print(combos(test))

If you do actually want 'a0' in the return value to be wrapped in an array, then you can replace ret.append(curcopy) with ret.append([curcopy]).

Upvotes: 1

Related Questions