Python - 获取n 1d数组中所有可能的元素总和

Question

给定一个整数n和一个数组a,我想返回一个数组,其中包含a的总和的所有可能值n次.

Example: n = 3, a = [1, 2, 3, 4, 5, 6]

Output: [3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18]

第一个元素是1 + 1 + 1,第二个是1 + 1 + 2等.

有没有优雅的方法呢？我已经尝试过循环,但由于事先不知道n,我不知道需要做多少循环.

提前致谢

Answer 1

生成所有可能的3元素组合,然后将它们相加:

from itertools import combinations_with_replacement

n = 3
li = [1, 2, 3, 4, 5, 6]

print([sum(comb) for comb in combinations_with_replacement(li, n)])

# [3, 4, 5, 6, 7, 8, 5, 6, 7, 8, 9, 7, 8, 9, 10, 9, 10, 11, 11, 12, 13, 6, 7, 8, 9, 10, 8, 9, 10, 11, 10, 11, 12, 12, 13, 14, 9, 10, 11, 12, 11, 12, 13, 13, 14, 15, 12, 13, 14, 14, 15, 16, 15, 16, 17, 18]

由于您似乎对唯一总和感兴趣,请使用集合:

print(set(sum(comb) for comb in combinations_with_replacement(li, n)))

# {3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18}

请注意,无法保证订购这些产品.如果你想要有序的输出明确它:

print(sorted(set(sum(comb) for comb in combinations_with_replacement(li, n))))