Question

在Python中使用numpy有几个优雅的例子可以生成所有组合的数组。例如答案：Using numpy to build an array of all combinations of two arrays。

现在假设有一个额外的约束，即所有数字的总和不能超过给定的常数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]])

原则上，来自https://stackoverflow.com/a/25655090/1479342的方法可用于生成没有约束的所有可能组合，然后选择总和小于K的组合子集。但是，该方法会产生比必要组合更多的组合，尤其是K与sum(maxRange)相比相对较小时。

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

Answer 1

我不知道numpy方法是什么，但这是一个相当干净的解决方案。让A为整数数组，让k为您输入的数字。

以空数组B开头;将数组B的总和保存在变量s中（最初设置为零）。应用以下程序：

如果数组s的总和B小于k，则（i）将其添加到集合，（ii）以及原始数组A中的每个元素，将该元素添加到B并更新s，（iii）将其从A中删除（iv））递归地应用程序; （iv）当呼叫返回时，将元素添加回A并更新s; 其他什么都不做。

这种自下而上的方法在早期修剪无效分支，只访问必要的子集（即几乎只有总和小于k的子集）。

Answer 2

<强>被修改

为了完整起见，我在这里添加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])

第一种方法是纯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)

经常性方法似乎比上述方法好得多：

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

经过短暂的思考，我能够更进一步。如果我们事先知道可能的分区数，我们可以一次分配足够的内存。（它与an already linked thread中的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

一些测试：

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

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

2 个答案: