python中优雅的多维循环

Question

我想用三个整数的平方和来表示哪个整数。在我的情况下，直到i，j，k都等于4。用于此目的的(不优雅的)代码如下：

nums = []
for i in range(4):
    for j in range(4):
        for k in range(4):
            nums.append(i*i + j*j + k*k)

对于这个问题，有没有更好的解决办法呢？理想情况下，对于任意多个变量(i，j，k，…，n)，我想要。

Answer 1

您可以使用itertools.product来完成以下操作：

from itertools import product

for i, j, k in product(range(4), repeat=3):
    # do something with i, j and k

如果事先不知道需要多少变量，就不必将product解压缩到i、j和k中。你可以把它作为：

for nums in product(range(4), repeat=3):
    # do something with nums

Answer 2

如果您想对更大的值执行此操作，我建议使用itertools

import itertools

def sumsq(x):
    return sum(y * y for y in x)

max_integer = 4
num_terms = 3

num = [sumsq(x) for x in itertools.combinations_with_replacement(range(1, max_integer + 1), num_terms)]

Answer 3

可以在一行中使用列表理解，如下所示：

l = [i*i + j*j + k*k for i in range(4) for j in range(4) for k in range(4)]