可变数量的相关嵌套循环
给出两个整数n
和d
,我想构造一个长度为d
的所有非负元组的列表,这些元组的总和为n
,包括所有排列.这类似于整数分区问题,但解决方案要简单得多.例如
d==3
:
Given two integers n
and d
, I would like to construct a list of all nonnegative tuples of length d
that sum up to n
, including all permutations. This is similar to the integer partitioning problem, but the solution is much simpler. For example for
d==3
:
[
[n-i-j, j, i]
for i in range(n+1)
for j in range(n-i+1)
]
可以很容易地将其扩展到更多维度,例如d==5
:
This can be extended to more dimensions quite easily, e.g., 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
,即嵌套循环的数量,一个变量,但是我不确定当时如何嵌套循环.
I would now like to make d
, i.e., the number of nested loops, a variable, but I'm not sure how to nest the loops then.
有任何提示吗?
递归:首先创建一个长度为d-1
的元组列表,该元组遍历所有ijk
,然后用另一列
Recursion to the rescue: First create a list of tuples of length d-1
which runs through all ijk
, then complete the list with another column 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)