CODEWITHSUNDEEP
pythonarraysnumpycombinationsscalar
Robi Roboter
Robi Roboter

Reputation: 43

Python Numpy: All combinations from array or scalar

I need from an array (or scalar) all possible combinations.

I have an Array np.array([0, 1, 0, 1]) or a scalar np.sum(np.array([0, 1, 0, 1])). Now i need all combinations such as:

[[2,0,0,0],
[1,1,0,0],
[1,0,1,0],
…
[0,0,1,1],
[0,0,0,2]]

The sum from an elements of all possible combinations can not exceed the sum of the scalar. I know the solution from "Using numpy to build an array of all combinations of two arrays“. But it does not help me for this problem. Have somebody any idea?

Upvotes: 1

Views: 3775

Answers (2)

Corralien
Corralien

Reputation: 120391

Several years later, maybe it could help someone.

The strategy is to generate all the combinations using itertools.product and keep right rows with numpy.where (which is faster).

import numpy as np
import itertools

a = np.array([0, 1, 0, 1])
l = len(a)
s = sum(a)

m = np.array(list(itertools.product(range(0, s+1), repeat=l)))
m = m[np.where(np.sum(m, axis=1) <= s)]

>>> m
array([[0, 0, 0, 0],
       [0, 0, 0, 1],
       [0, 0, 0, 2],
       [0, 0, 1, 0],
       [0, 0, 1, 1],
       [0, 0, 2, 0],
       [0, 1, 0, 0],
       [0, 1, 0, 1],
       [0, 1, 1, 0],
       [0, 2, 0, 0],
       [1, 0, 0, 0],
       [1, 0, 0, 1],
       [1, 0, 1, 0],
       [1, 1, 0, 0],
       [2, 0, 0, 0]])

Upvotes: 0

or1426
or1426

Reputation: 949

Here is a very simple (and inefficient) way of doing what you want for the scalar case. For the array case just get the numbers n and m from the sum and length of the array. I'm sure its possible to do significantly better than this with some additional thinking but unless your numbers get significantly larger then this will suffice:

import itertools
import numpy as np

n=2 #number of elements to distribute
m=4 #length of arrays to distribute in
arrays = []
for combination in itertools.combinations_with_replacement(range(m), n):
    arrays.append(np.zeros(m,dtype=int))
    for index in combination:
        arrays[-1][index] += 1

for array in arrays:
    print(array)

Output:

[2 0 0 0]
[1 1 0 0]
[1 0 1 0]
[1 0 0 1]
[0 2 0 0]
[0 1 1 0]
[0 1 0 1]
[0 0 2 0]
[0 0 1 1]
[0 0 0 2]

Upvotes: 3

Related Questions