python combinations of split of list

user_17346587 0
Viewed 52

Given two input arguments:

x: <list> like [0, 1, 2, 3]

splits: <list> of int representing the number and the length of the splits, like [1, 2, 1]. It means that the list must be split in 3, the 1st one contains 1 element, the 2nd contains 2 elements and the 3rd split contains 1 elements

I need a function that returns all the possible combinations of the splits, eg:

def get_combos(x, splits): ...

>>> get_combos([0, 1, 2, 3], [1, 2, 1])

[(0,), (1, 2), (3,)]
[(0,), (1, 3), (2,)]
[(0,), (2, 3), (1,)]
[(1,), (0, 2), (3,)]
[(1,), (0, 3), (2,)]
[(1,), (2, 3), (0,)]
[(2,), (0, 1), (3,)]
[(2,), (0, 3), (1,)]
[(2,), (1, 3), (0,)]
[(3,), (0, 1), (2,)]
[(3,), (0, 2), (1,)]
[(3,), (1, 2), (0,)]

Input will always meet these conditions:

  • min(splits) >= 1 --> There is always an element in the split
  • sum(splits) == len(x) --> All the elements in x are taken

What is the best way to do it recursively?

combinations combinatorics recursion python
1 Answers

You can use a recursive generator function:

def combos(a, b, c = []):
   if not a or not b:
      yield [*map(tuple, c)]
   else:
      for x, y in enumerate(a):
         if not c or len(c[-1]) == b[0]:
            yield from combos(a[:x]+a[x+1:], b[1:] if c else b, c+[[y]])
         else:
            yield from combos(a[:x]+a[x+1:], b, [*c[:-1], c[-1]+[y]])

print(list(combos([0, 1, 2, 3], [1, 2, 1])))

Output:

[[(0,), (1, 2), (3,)], [(0,), (1, 3), (2,)], [(0,), (2, 1), (3,)], [(0,), (2, 3), (1,)], [(0,), (3, 1), (2,)], [(0,), (3, 2), (1,)], [(1,), (0, 2), (3,)], [(1,), (0, 3), (2,)], [(1,), (2, 0), (3,)], [(1,), (2, 3), (0,)], [(1,), (3, 0), (2,)], [(1,), (3, 2), (0,)], [(2,), (0, 1), (3,)], [(2,), (0, 3), (1,)], [(2,), (1, 0), (3,)], [(2,), (1, 3), (0,)], [(2,), (3, 0), (1,)], [(2,), (3, 1), (0,)], [(3,), (0, 1), (2,)], [(3,), (0, 2), (1,)], [(3,), (1, 0), (2,)], [(3,), (1, 2), (0,)], [(3,), (2, 0), (1,)], [(3,), (2, 1), (0,)]]
user_73267380
Related