可变数量的从属嵌套循环

时间:2017-07-27 10:24:41

标签: python nested-loops integer-partition

给定两个整数nd,我想构建一个长度为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,即嵌套循环的数量,一个变量,但我不知道如何嵌套循环。

任何提示?

1 个答案:

答案 0 :(得分:5)

拯救的递归:首先创建一个长度为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)