3

这个问题类似于我几个月前的一个问题：Generating a numpy array with all combination of numbers that sum than a given number。在那个问题中，我想生成所有总和最多为常数的数字，因为每个元素都有一个最大值。

这一次，我想计算总和恰好等于该常数的所有排列。这可以看作是计算整数分区的唯一排列，其中每个元素都有一定的最大值。最终结果应该存储在一个 numpy 数组中。

使用生成器，one liner 实现了我们想要的：

import numpy as np
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, 1, 2],
       [0, 2, 1],
       [0, 3, 0],
       [1, 0, 2],
       [1, 1, 1],
       [1, 2, 0]])

K=20当和时，我的性能很慢maxRange = [20]*6。排列的数量被限制为 53130，但它已经需要 20 秒。我的直觉告诉我，这应该花费不到一秒钟的时间。

有没有人有更快的解决方案？我无法修改我之前的问题的解决方案来解决这个问题，因为我不知道如何切断不再可能完全加起来的排列K。

我不介意使用@jitnumba 运算符的解决方案......只要它们比我现在拥有的更快！

提前致谢。

4

1 回答 1

1

我不得不考虑很长时间，但我设法修改了生成一个numpy数组的解决方案，其中所有数字组合的总和小于这个问题的给定数字：

对于分区的数量，我们的想法是计算一个数组，该数组feasible_range指定我们在某个阶段至少总共需要多少才能达到max_sum。例如，如果我们想总共达到 3 个 和max_range[0] == 1，那么我们需要至少有 2 个才能从最后一个元素开始。这个数组来自一个累积和：

feasible_range = np.maximum(max_sum - np.append(np.array([0]),np.cumsum(max_range)[:-1]),0)

现在我们可以像以前一样计算分区的数量，方法是将永远不会导致可行分区的元素设置为 0。

def number_of_partitions(max_range, max_sum):
    M = max_sum + 1
    N = len(max_range) 
    arr = np.zeros(shape=(M,N), dtype = int)    
    feasible_range = max_sum - np.append(np.array([0]),np.cumsum(max_range)[:-1])
    feasible_range = np.maximum(feasible_range,0)     

    arr[:,-1] = np.where(np.arange(M) <= min(max_range[-1], max_sum), 1, 0) 
    arr[:feasible_range[-1],-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] 
        arr[:feasible_range[i],i]=0  #Set options that will never add up to max_sum at 0.
    return arr.sum(axis = 0),feasible_range

配分函数也有与之前类似的解释。

def partition(max_range, max_sum, out = None, n_part = None,feasible_range=None):
    #Gives all possible partitions of the sets 0,...,max_range[i] that sum up to  max_sum. 
    if out is None:
        max_range = np.asarray(max_range, dtype = int).ravel()
        n_part,feasible_range = 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(feasible_range[0],min(max_range[0],max_sum) + 1, dtype = int).reshape(-1,1)
        return out

    #Copy is needed since otherwise we overwrite some values of P.
    P = partition(max_range[1:], max_sum, out=out[:n_part[1],1:], n_part = n_part[1:],feasible_range=feasible_range[1:]).copy()      


    S = max_sum - P.sum(axis = 1) #The remaining space in the partition 
    offset, sz  = 0, 0
    for i in range(max_range[0]+1):
        #select indices for which there is remaining space
        #do this only if adding i brings us within the feasible_range.
        ind, = np.where(np.logical_and(S-i>=0,S-i <= max_sum-feasible_range[0])) 
        offset, sz = offset + sz, ind.size
        out[offset:offset+sz, 0] = i
        out[offset:offset+sz, 1:] = P[ind]
    return out

对于K=20和maxRange = [20]*6，partition(maxRange,K)需要 13 毫秒，而首先需要 18.5 秒。

我不太喜欢必须复制的部分；这可能可以通过颠倒顺序来避免。不过现在速度已经够好了。

于 2016-09-14T10:22:07.910 回答