Python - 获取 n 个一维数组中所有可能的元素总和

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

生成所有可能的三元素组合，然后对它们求和:

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