在這種情況下，itertools可以用於數量不確定的尺寸嗎？

Question

這是我的代碼，使用dim = 3，但是我希望它適用於任何維度，而無需手動編輯代碼。

我希望能夠最終在3到20之間更改尺寸，而無需手動添加for循環。 我一直在查看itertools，但不知道如何從itertools.product()創建的元組中選擇正確的值以平方並為if語句itertools.product() 。

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

作為獎勵，您擺脫了unpythonic range(len())構造。

當您擁有向量載體時，我已經做了一個更一般的案例。 讀起來有點困難，因為在for循環中沒有完成拆包。

平方和使用做sum在1個指數（值），如果條件匹配，我們循環，直到我們找到了“更深層次”列表中的值設置為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是一個numpy ndarray ，我們可以通過傳遞索引列表來簡化“設置為1”循環：

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