Search code examples
pythonpython-3.xlistnestedpermutation

Include permutations of nested lists


I have the following function to get the number of permutations (without repeating elements) of a list:

import itertools

def permutations_without_repetition(samples, size):
    return list(itertools.permutations(samples, size))

Which works fine for me as long as the list I provide doesn't contain nested lists. Itertools treats nested elements just as one whole element, and doesn't bother also generating permutations containing differently-ordered nested lists.

If I run permutations_without_repetition([[1, 2], 3], 2), the only results I get are:

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

I know this is the expected behaviour, but I would like the result to be:

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

What's the easiest way to also return permutations that contain permutations of nested lists, to produce the result above?


Solution

  • You can use a recursive generator function:

    def combos(d, c = []):
       if not isinstance(d, list):
          yield d
       elif not d:
          yield c
       else:
          for i, a in enumerate(d):
             for k in combos(a, c = []):
                 yield from combos(d[:i]+d[i+1:], c+[k])
    
    print(list(combos([[1, 2], 3])))
    

    Output:

    [[[1, 2], 3], [[2, 1], 3], [3, [1, 2]], [3, [2, 1]]]
    

    Shorter solution using itertools:

    import itertools as it
    def combos(d):
       if not isinstance(d, list):
          yield d
       else:
          for i in it.permutations(d):
             yield from map(list, it.product(*[combos(j) for j in i]))
    
    print(list(combos([[1, 2], 3])))
    

    Output:

    [[[1, 2], 3], [[2, 1], 3], [3, [1, 2]], [3, [2, 1]]]