在这种情况下，itertools是否可以用于未指定数量的维度？

Question

这是我的代码，使用dim=3，但是我希望它可以在任何维度上工作，而不必手动编辑代码

我想能够改变3和20之间的维度最终没有手动广告为循环。 我在看itertools，但不知道如何从itertools.product()创建的元组中选择正确的值来计算if语句的平方和总和

arrayshape = (width * 2 + 1,) * dim
funcspace = np.zeros(shape=arrayshape, dtype='b')

x1 = list(range(-int(width), int(width + 1)))
x2 = x1
x3 = x1

for i in range(len(x1)):
    for j in range(len(x2)):
        for k in range(len(x3)):
            if round(np.sqrt(x1[i] ** 2 + x2[j] ** 2 + x3[k] ** 2)) in ranges:
                funcspace[i][j][k] = 1

Answer 1

您可以在向量的enumerate上使用product，这将产生值和索引：

for ((i,v1),(j,v2),(k,v3)) in itertools.product(enumerate(x1),enumerate(x2),enumerate(x3)):
    if round(np.sqrt(v1**2+v2**2+v3**2)) in ranges:
            funcspace[i][j][k]=1

作为奖励，您摆脱了非同音的range(len())构造

我做了一个更一般的例子，当你有一个向量向量。读起来有点难，因为解包不是在for循环中完成的

平方和是在1索引（值）上使用sum来完成的，如果条件匹配，我们将循环，直到找到“更深层”的列表来将值设置为1

for t in itertools.product(*(enumerate(x) for x in x_list)):
    # compute the squared sum of values
    sqsum = sum(v[1]**2 for v in t)
    if round(sqsum) in ranges:
        # traverse the dimensions except the last one
        deeper_list = funcspace
        for i in range(len(t)-1):
            deeper_list = deeper_list[t[i][0]]
        # set the flag using the last dimension list
        deeper_list[t[-1][0]] = 1

如注释中所述，由于x1似乎重复出现，您可以用以下语句替换第一个语句：

for t in itertools.product(enumerate(x1), repeat=dim):

另一条评论指出，由于funcspace是一个numpyndarray，我们可以通过传递索引列表来简化“set to 1”循环：

funcspace[[x[0] for x in t]] = 1