Search code examples
pythonlistpython-itertools

Get combination of sublists with length n

Say I have a list lst = [[[24, 25], [25, 26], [25, 35]], [[24, 25, 26], [26, 36, 46]]]. I want to calculate all combinations with all elements from other sublists, but not the same sublist. In this example it would be:

[[24, 25], [24, 25, 26]],
[[25, 26], [24, 25, 26]],
[[25, 35], [24, 25, 26]],
[[24, 25], [26, 36, 46]],
[[25, 26], [26, 36, 46]],
[[25, 35], [26, 36, 46]]

This should however also be able to produce combinations of more than 2 elements, i.e. lst = [[[1]], [[2], [3]], [[4], [5]]]; (length=len(lst) -> 3):

[[1], [2], [4]],
[[1], [2], [5]],
[[1], [3], [4]],
[[1], [3], [5]]

I tried with itertools.combinations(*lst) but that only netted tuples of length 2. How can I achieve combinations of length n with the limitations outlined above?

Solution

  • Use itertools.product to generate the cartesian product over two or more iterables:

    from itertools import product

top = [[[24, 25], [25, 26], [25, 35]], [[24, 25, 26], [26, 36, 46]]]

for combo in product(*top):
  print(combo)

    Outputs:

    ([24, 25], [24, 25, 26])
([24, 25], [26, 36, 46])
([25, 26], [24, 25, 26])
([25, 26], [26, 36, 46])
([25, 35], [24, 25, 26])
([25, 35], [26, 36, 46])