python - 生成总和为参数 N 的列表的乘积

Question

我正在寻找一种快速方法来生成值从 0 到(包括)N 的乘积，总和为 N。对于更简单的问题，这看起来像

N = 10
output = []
for i in xrange(N+1):
    for j in xrange(N+1-i):
        for k in xrange(N+1-i-j):
           output.append([i,j,k])   

所以，在这种情况下，输出应该只包含长度为 3 的列表。类似于

[v for v in itertools.product(*[range(N+1) for i in range(3)]) if sum(v) == N)]

编辑:正如 Joran 指出的那样，上面的行没有利用第三个元素将始终是“其余”的事实，即 N-i-j。

希望有一种方法可以使用 itertools 执行此操作，而无需评估每种可能产品的结果。目前，我有一个递归函数，除了一些效率问题(例如经常调用 len)，对于较大的 max_depth 需要相当长的时间:

def recursive_expand(input_list, max_sum, max_depth):
    ''' Creates list of lists of integers, where elements of each 
        sub_list sum to max_sum.
    '''
    # If current elements already sum to max_sum
    if max_sum == 0:
        return [input_list + [0] * (max_depth - len(input_list))]
    # If the list can only contain one more element
    elif len(input_list) == max_depth - 1:
        return [input_list + [max_sum]]

    output_lists = []
    for n in xrange(max_sum + 1):
        output_lists.extend( \
                 recursive_expand(input_list + [n], max_sum - n,max_depth))
    return output_lists

>>> foo = recursive_expand([], 2, 3)
>>> print np.array(foo)
[[0 0 2]
 [0 1 1]
 [0 2 0]
 [1 0 1]
 [1 1 0]
 [2 0 0]]

编辑 2:归功于 jonrsharpe，此函数生成长度为 k 的元组，其中包含从 start 到 n 的整数，求和为 n，其中顺序很重要。如果顺序无关紧要(因为这样效率更高)，请参阅他的回答。

def sets_2(n, k, start):
    for x in xrange(start, n + 1):
        if k == 2:
            yield x, n-x
        elif n-x == 0:
            yield (x,) + (0,) * (k-1)
        else:
            for tup in sets(n-x, k-1, x):
                yield (x,) + tup

而且，发现整数分区(之后可以用零填充)的最有效实现可以在这里找到:http://homepages.ed.ac.uk/jkellehe/partitions.php

Answer 1

如果您对顺序不感兴趣(即您总是想要 x <= y <= z ，只有唯一的整数集)，这应该是相当有效的，通过:

1. 限制range您搜索合适的x和 y ;
2. 正在计算 z来自 x和 y ，而不是再次循环；和
3. yield计算值以创建生成器。

代码:

def sets(n):
    for x in xrange(1, (n // 3) + 1):
       for y in xrange(x, ((n - x) // 2) + 1):
           yield x, y, n - (x + y)

例如:

list(sets(10)) == [(1, 1, 8), (1, 2, 7), (1, 3, 6), (1, 4, 5), 
                   (2, 2, 6), (2, 3, 5), (2, 4, 4), (3, 3, 4)]

对于任意设置长度k ，你可以使用递归:

def sets_2(n, k=3, s=1):
    for x in range(s, (n // k) + 1):
        if k == 2:
            yield x, n - x
        else:
            for s in sets_2(n - x, k - 1, x):
                yield (x, ) + s

例如:

list(sets_2(10, 2)) == [(1, 9), (2, 8), (3, 7), (4, 6), (5, 5)]

list(sets_2(10, 4)) == [(1, 1, 1, 7), (1, 1, 2, 6), (1, 1, 3, 5), 
                        (1, 1, 4, 4), (1, 2, 2, 5), (1, 2, 3, 4), 
                        (1, 3, 3, 3), (2, 2, 2, 4), (2, 2, 3, 3)]

要包含零，请设置 s=0 :

list(sets_2(10, 3, 0)) == [(0, 0, 10), (0, 1, 9), (0, 2, 8), (0, 3, 7), 
                           (0, 4, 6), (0, 5, 5), (1, 1, 8), (1, 2, 7), 
                           (1, 3, 6), (1, 4, 5), (2, 2, 6), (2, 3, 5), 
                           (2, 4, 4), (3, 3, 4)]