python - 可变数量的依赖嵌套循环

Question

给定两个整数 n 和 d，我想构建一个列表，其中包含长度为 d 的所有非负元组，其总和为 n，包括所有排列。这类似于 integer partitioning problem ，但解决方案要简单得多。例如对于 d==3:

[
    [n-i-j, j, i]
    for i in range(n+1)
    for j in range(n-i+1)
]

这可以很容易地扩展到更多维度，例如，d==5:

[
    [n-i-j-k-l, l, k, j, i]
    for i in range(n+1)
    for j in range(n-i+1)
    for k in range(n-i-j+1)
    for l in range(n-i-j-l+1)
]

我现在想将 d(即嵌套循环的数量)设置为一个变量，但我不确定如何嵌套循环。

有什么提示吗？

Answer 1

递归救援:首先创建一个长度为 d-1 的元组列表，它遍历所有 ijk，然后用另一列 n 完成该列表-sum(ijk).

def partition(n, d, depth=0):
    if d == depth:
        return [[]]
    return [
        item + [i]
        for i in range(n+1)
        for item in partition(n-i, d, depth=depth+1)
        ]


# extend with n-sum(entries)
n = 5
d = 3
lst = [[n-sum(p)] + p for p in partition(n, d-1)]

print(lst)