combination of numpy array satisfying condition

Question

I am trying to produce all combination of numpy array that satisfy a condition efficiently my code now looks like this

import numpy as np
import itertools

a = np.array([1,11,12,13])
a = np.tile(a,(13,1))
a = a.flatten()

for c in itertools.combinations(a,4):
  if np.sum(c)==21:
    print(c)

a = np.tile(a,4) is sufficient for this problem. Are you looking at a general solution ? — B. M., Apr 18 '16 at 11:36

score 0 · Accepted Answer · edited May 23 '17 at 12:09

If you only care about unique combinations (and there are only 256 of them), you can use itertools.product:

version_1 = np.vstack(list(sorted({tuple(row) for row in list(itertools.combinations(a, 4))})))  # unique combinations, your way
version_2 = np.array(list(itertools.product((1, 11, 12, 13), repeat=4)))  # same result, but faster
assert (version_1 == version_2).all()

I'm using this answer to get the unique elements of a Numpy array.

So the final answer would be:

import itertools, numpy as np
a = np.array(list(itertools.product((1, 11, 12, 13), repeat=4)))
for arr in a[a.sum(axis=1) == 21]:
    print(arr)

combination of numpy array satisfying condition

1 Answers1