Creating a numpy array with all combinations of numbers whose sum is less than a given number

Question

There are some elegant examples of using numpy in Python to generate arrays of all combinations. For example, the answer here is: Using numpy to build an array of all combinations of two arrays .

Now suppose that there is an additional restriction, namely, the sum of all numbers cannot contain more than a given constant K Using the generator and itertools.product , for an example with K=3 , where we want a combination of three variables with ranges 0-1.0-3 and 0-2, we can do this:

 from itertools import product K = 3 maxRange = np.array([1,3,2]) states = np.array([i for i in product(*(range(i+1) for i in maxRange)) if sum(i)<=K]) 

which returns

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

Basically, the approach of https://stackoverflow.com/a/377832/ can be used to create all possible combinations without restriction, and then select a subset of combinations whose sum is less than K However, this approach generates many more combinations than necessary, especially if K relatively small compared to sum(maxRange) .

There must be a way to do this faster and with less memory use. How can this be achieved using a vectorized approach (e.g. using np.indices )?

Answer 1

Edited

• For completeness, I add the OP code here:

   def partition0(max_range, S): K = len(max_range) return np.array([i for i in itertools.product(*(range(i+1) for i in max_range)) if sum(i)<=S]) 

• The first approach is pure np.indices . This is fast for small input, but consumes a lot of memory (the OP has already pointed out that this is not what he had in mind).

   def partition1(max_range, S): max_range = np.asarray(max_range, dtype = int) a = np.indices(max_range + 1) b = a.sum(axis = 0) <= S return (a[:,b].T) 

• The repetitive approach seems a lot better than the above:

   def partition2(max_range, max_sum): max_range = np.asarray(max_range, dtype = int).ravel() if(max_range.size == 1): return np.arange(min(max_range[0],max_sum) + 1, dtype = int).reshape(-1,1) P = partition2(max_range[1:], max_sum) # S[i] is the largest summand we can place in front of P[i] S = np.minimum(max_sum - P.sum(axis = 1), max_range[0]) offset, sz = 0, S.size out = np.empty(shape = (sz + S.sum(), P.shape[1]+1), dtype = int) out[:sz,0] = 0 out[:sz,1:] = P for i in range(1, max_range[0]+1): ind, = np.nonzero(S) offset, sz = offset + sz, ind.size out[offset:offset+sz, 0] = i out[offset:offset+sz, 1:] = P[ind] S[ind] -= 1 return out 

• After a short thought, I was able to take this a little further. If we know in advance the number of possible partitions, we can allocate enough memory right away. (It looks a bit like cartesian in an already connected thread .)

  First, we need a function that takes partitions into account.

   def number_of_partitions(max_range, max_sum): ''' Returns an array arr of the same shape as max_range, where arr[j] = number of admissible partitions for j summands bounded by max_range[j:] and with sum <= max_sum ''' M = max_sum + 1 N = len(max_range) arr = np.zeros(shape=(M,N), dtype = int) arr[:,-1] = np.where(np.arange(M) <= min(max_range[-1], max_sum), 1, 0) for i in range(N-2,-1,-1): for j in range(max_range[i]+1): arr[j:,i] += arr[:Mj,i+1] return arr.sum(axis = 0) 

  The main function:

   def partition3(max_range, max_sum, out = None, n_part = None): if out is None: max_range = np.asarray(max_range, dtype = int).ravel() n_part = number_of_partitions(max_range, max_sum) out = np.zeros(shape = (n_part[0], max_range.size), dtype = int) if(max_range.size == 1): out[:] = np.arange(min(max_range[0],max_sum) + 1, dtype = int).reshape(-1,1) return out P = partition3(max_range[1:], max_sum, out=out[:n_part[1],1:], n_part = n_part[1:]) # P is now a useful reference S = np.minimum(max_sum - P.sum(axis = 1), max_range[0]) offset, sz = 0, S.size out[:sz,0] = 0 for i in range(1, max_range[0]+1): ind, = np.nonzero(S) offset, sz = offset + sz, ind.size out[offset:offset+sz, 0] = i out[offset:offset+sz, 1:] = P[ind] S[ind] -= 1 return out 

• Some tests:

   max_range = [3, 4, 6, 3, 4, 6, 3, 4, 6] for f in [partition0, partition1, partition2, partition3]: print(f.__name__ + ':') for max_sum in [5, 15, 25]: print('Sum %2d: ' % max_sum, end = '') %timeit f(max_range, max_sum) print() partition0: Sum 5: 1 loops, best of 3: 859 ms per loop Sum 15: 1 loops, best of 3: 1.39 s per loop Sum 25: 1 loops, best of 3: 3.18 s per loop partition1: Sum 5: 10 loops, best of 3: 176 ms per loop Sum 15: 1 loops, best of 3: 224 ms per loop Sum 25: 1 loops, best of 3: 403 ms per loop partition2: Sum 5: 1000 loops, best of 3: 809 µs per loop Sum 15: 10 loops, best of 3: 62.5 ms per loop Sum 25: 1 loops, best of 3: 262 ms per loop partition3: Sum 5: 1000 loops, best of 3: 853 µs per loop Sum 15: 10 loops, best of 3: 59.1 ms per loop Sum 25: 1 loops, best of 3: 249 ms per loop 

  And something more:

   %timeit partition0([3,6] * 5, 20) 1 loops, best of 3: 11.9 s per loop %timeit partition1([3,6] * 5, 20) The slowest run took 12.68 times longer than the fastest. This could mean that an intermediate result is being cached 1 loops, best of 3: 2.33 s per loop # MemoryError in another test %timeit partition2([3,6] * 5, 20) 1 loops, best of 3: 877 ms per loop %timeit partition3([3,6] * 5, 20) 1 loops, best of 3: 739 ms per loop 

Answer 2

I do not know what numpy , but here is a fairly clean solution. Let A be an array of integers, and k number you specify as input.

Start with an empty array B ; save the sum of the array B in the variable s (initially set to zero). Apply the following procedure:

• if the sum s array B less than k , and then (i) add it to the collection, (ii) and for each element from the original array A add this element to B and update s , (iii) remove it from A and (iv ) apply the procedure recursively; (iv) when the call returns, add the item back to A and update s ; else do nothing.

This bottom-up approach first cuts out invalid branches and only visits the necessary subsets (i.e. almost only subsets whose sum is less than k ).