生成与参数N相加的列表的乘积

时间:2014-02-11 18:11:53

标签: python

我正在寻找一种快速的方法来生成值为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)]
编辑:正如约兰所指出的那样,上面的几行并没有利用第三个元素永远是“休息”的事实,即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的元组,其中包含从startn的整数,汇总到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

2 个答案:

答案 0 :(得分:2)

如果你对订单不感兴趣(即你总是想要x <= y <= z,只有唯一的整数集),这应该是合理有效的,通过:

  1. 限制range您搜索适当的xy;
  2. zx计算y,而不是再次循环;和
  3. yield创建生成器的值。
  4. 代码:

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

答案 1 :(得分:1)

好的我想在这里走出困境并猜测你想要3个值加起来为某个神奇值N,同时还要满足一些额外的约束

你不需要计算所有3个值

a + b +c = N  ==> c = N - a -b

你可以计算2个值并解决第3个

for a in range(N): #unless the other 2 values can be 0 you dont really need to go all the way to n
    for b in range(a,N):  #this assumes that b >= a
         c = N - a - b
         if some_constraint(a,b,c,N):
                my_list.append([a,b,c])
说实话,我不知道这是否能回答你的问题,但看起来这就是你问题的方向......