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