使用特定總和迭代列表

Question

我想迭代所有長度為n的列表，其元素總和為2.你怎么能有效地做到這一點？ 對於n = 10這是一種非常低效的方法。 最終我想為'n> 25'這樣做。

n = 10
for L in itertools.product([-1,1], repeat = n):
    if (sum(L) == 2):
        print L #Do something with L

Answer 1

如果你有2個+1而不是-1，你只能得到2的解，所以對於n == 24

a_solution = [-1,]*11 + [1,]*13  

現在你可以使用itertools.permutations來獲得它的每一個排列

for L in itertools.permutations(a_solution): print L

使用itertools.combinations來消除重復可能會更快

for indices in itertools.combinations(range(24),11):
    a = numpy.ones(24)
    a[list(indices)] = -1
    print a

注意到你得到2列表必須是一個均勻的長度

Answer 2

一種方法是在剩余元素無法彌補目標總和時停止產品遞歸。

具體來說，此方法將您的itertools.product(...,repeat)分成一個遞歸生成器，它根據當前列表元素的值更新目標總和，並在遞歸之前檢查結果目標是否可實現進一步：

def generate_lists(target, n):
  if(n <= 0):
    yield []
    return
  if(target > n or target < -n):
    return
  for element in [-1,1]:
    for list in generate_lists(target-element, n-1):
      yield list+[element]