生成一个numpy数组,其中包含总和小于给定数字的所有数字组合

Question

在Python中使用numpy生成所有组合的数组有几个优雅的例子.例如答案:使用numpy构建两个数组的所有组合的数组.

现在假设存在一个额外的约束,即所有数字的总和不能超过给定的常数K.使用生成器itertools.product,例如K=3我们想要三个变量的组合,范围为0-1,0-3和0-2,我们可以这样做:

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])

返回

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]])

原则上,来自/sf/answers/1795856331/的方法可用于生成所有可能的组合而不受约束,然后选择总和小于的组合子集K.然而,这种方法产生了比必要组合更多的组合,特别是如果K相对较小的组合sum(maxRange).

必须有一种方法可以更快地执行此操作并降低内存使用率.如何使用矢量化方法(例如使用np.indices)实现这一目标？

Answer 1

已编辑

1. 为了完整起见，我在此处添加OP的代码：

   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])
   

2. 第一种方法是纯净的np.indices。对于少量输入而言，它速度很快，但会占用大量内存（OP已经指出这不是他的意思）。

   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)
   

3. 循环方法似乎比上述方法要好得多：

   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
   

4. 经过短暂的思考，我可以将其进一步介绍。如果我们事先知道可能的分区数，则可以一次分配足够的内存。（这有点类似于cartesian在一个已经连接线。）

   首先，我们需要一个计算分区的函数。

   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[:M-j,i+1] 
       return arr.sum(axis = 0)
   

   主要功能：

   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
   

5. 一些测试：

   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
   

   还有更大的东西：

   %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