python - Get all permutations of a numpy array

16
votes

I have a numpy array [0, 1, 1, 2, 2, 0, 1, ...] which only contains the numbers 0-k. I would like to create a new array that contains the n possible arrays of permutations of 0-k. A small example with k=2 and n=6:

a = [0, 1, 0, 2]
permute(a)
result = [[0, 1, 0, 2]
          [0, 2, 0, 1]
          [1, 0, 1, 2]
          [2, 1, 2, 0]
          [1, 2, 1, 0]
          [2, 0, 2, 1]]

Does anyone have any ideas/solutions as to how one could achieve this?

2
pythonnumpy
This is a duplicate plus a google away. – kabanus
@kabanus OP seems to want distinguishable permutations -- which that link doesn't give. – John Coleman

2 Answers

36
votes

Your a is what combinatorists call a multiset. The sympy library has various routines for working with them.

>>> from sympy.utilities.iterables import multiset_permutations
>>> import numpy as np
>>> a = np.array([0, 1, 0, 2])
>>> for p in multiset_permutations(a):
...     p
...     
[0, 0, 1, 2]
[0, 0, 2, 1]
[0, 1, 0, 2]
[0, 1, 2, 0]
[0, 2, 0, 1]
[0, 2, 1, 0]
[1, 0, 0, 2]
[1, 0, 2, 0]
[1, 2, 0, 0]
[2, 0, 0, 1]
[2, 0, 1, 0]
[2, 1, 0, 0]
25
votes

if your permutations fit in the memory, you could store them in a set and thus only get the distinguishable permutations.

from itertools import permutations

a = [0, 1, 0, 2]

perms = set(permutations(a))