动态嵌套循环 Python

问题描述

这是我需要的数组：

N = 6
A = [[x,y,z] for x in range(N+1) for y in range(N+1) for z in range(N+1) if x+y+z== N]

是否有其他方法可以仅指定变量 N 和 3 而不是 x,z？
我尝试了 []*3，但无法获得所需的输出。

解决方法

您正在寻找iterttols.product：

from itertools import product

N = 6
A = [(x,y,z) for x in range(N+1) for y in range(N+1) for z in range(N+1) if x+y+z== N]
B = [tup for tup in product(range(N+1),repeat=3) if sum(tup) == N]
print(A,B,A == B,sep='\n')

给出：

[(0,6),(0,1,5),2,4),3,3),4,2),5,1),6,0),(1,(2,(3,(4,(5,(6,0)]
[(0,0)]
True

,

使用itertools.combinations_with_replacement：

from itertools import combinations_with_replacement
n = 6
k = 3

# If you want a list of tuples:
lst = [item for item in list(combinations_with_replacement(range(n),k)) if sum(item) == n]
print(lst)
# [(0,2)]

# If you want a list of lists:
lst = [list(item) for item in list(combinations_with_replacement(range(n),k)) if sum(item) == n]
print(lst)
# [[0,5],[0,4],3],[1,[2,2]]