Question

我想做这样的事情：

    if dim==2:
        a,b=grid_shape
        for i in range(a):
            for j in range(b):
                A[i,j] = ...things...

其中dim只是我的元组grid_shape中元素的数量。 A是一个维度dim的数组。有没有办法在没有特定尺寸的情况下做到这一点？无需编写像

这样的丑陋代码

    if dim==2:
        a,b=grid_shape
        for i in range(a):
            for j in range(b):
                A[i,j] = ...things...
    if dim==3:
        a,b,c=grid_shape
        for i in range(a):
            for j in range(b):
                for k in range(c):
                    A[i,j,k] = ...things...

Answer 1

使用itertools，您可以这样做：

for index in itertools.product(*(range(x) for x in grid_shape)):
    A[index] = ...things...

这取决于几个技巧。首先，itertools.product()是一个从迭代中生成元组的函数。

for i in range(a):
    for j in range(b):
        index = i,j
        do_something_with(index)

可以缩减为

for index in itertools.product(range(a),range(b)):
    do_something_with(index)

这适用于itertools.product()的任意数量的参数，因此您可以有效地创建任意深度的嵌套循环。

另一个技巧是将网格形状转换为itertools.product的参数：

(range(x) for x in grid_shape)

相当于

(range(grid_shape[0]),range(grid_shape[1]),...)

也就是说，它是每个grid_shape维度的范围元组。使用*然后将其扩展为参数。

itertools.product(*(range(x1),range(x2),...))

相当于

itertools.product(range(x1),range(x2),...)

此外，由于A[i,j,k]相当于A[(i,j,k)]，我们可以直接使用A[index]。

正如DSM所指出的，因为你正在使用numpy，你可以减少

itertools.product(*(for range(x) for x in grid_shape))

到

numpy.ndindex(grid_shape)

所以最终循环变为

for index in numpy.ndindex(grid_shape):
    A[index] = ...things...

Answer 2

你可以通过在最后一个变量前放一个星号来捕捉元组的其余部分，并通过在它周围加上括号来创建一个数组。

>>> tupl = ((1, 2), 3, 4, 5, 6)
>>> a, *b = tupl
>>> a
(1, 2)
>>> b
[3, 4, 5, 6]
>>>

然后你可以循环浏览b。所以它看起来像

a,*b=grid_shape
for i in a:
    for j in range(i):
        for k in b:
            for l in range(k):
                A[j, l] = ...things...

如何解压缩元组以进行循环而不是特定于维度

2 个答案: