如何在 python 中重写这个嵌套的 for 循环？

Question

我有一个嵌套的 for 循环，我正在对其进行一些操作，但是运行大约需要 20 分钟。希望减少上墙时间。这只是一个制作可重现的简单示例，但我如何重写这个嵌套的 for 循环以使其更有效？

我尝试在python中使用zip函数，但它没有像嵌套循环那样打印出i,j的顺序，这对我的计算是必要的.

这个问题类似于这个 Whosebug 问题，但我正在努力重现答案： Automatically nested for loops in python

array = np.random.rand(3,4,10)

x_vals = np.array([22,150,68,41]) #size of 4 (same as size as 2nd axis of array above)

new_arr = np.zeros((array.shape))

for i in range(3): 
     for j in range(4): 
            print(i,j)
            new_arr[i,:,j] = array[i,:,j]*x_vals
            
0 0
0 1
0 2
0 3
1 0
1 1
1 2
1 3
2 0
2 1
2 2
2 3

我试过了：

for i,j in zip(range(3),range(4)):
    print(i,j) # output is i,j is not the same as it is for the nested loop above

0 0
1 1
2 2

我在想也许函数 enumerate 也可以工作，但我也收到错误消息：

for idx,i in enumerate(range(3),range(4)):
    print(idx,i)

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
/tmp/ipykernel_2101160/1601005143.py in <module>
----> 1 for idx,i in enumerate(range(3),range(4)):
      2     print(idx,i)

TypeError: 'range' object cannot be interpreted as an integer

关于如何矢量化或加速嵌套循环的任何想法？

Answer 1

您可以按如下方式使用 itertools.product 以避免嵌套 fors

import itertools
for i,j in itertools.product(range(3),range(4)):
    print(i,j)

产出

不过对所需时间影响不大。你的提及

order of i,j(...)which is necessary for my calculation.

建议您需要先前计算的效果，即关于不能被破坏的顺序的某些规则。您需要考虑如何在不违反规则的情况下进行更改，以及它是否会提供所需的速度提升。

Answer 2

您可以使用 itertools.product，相当于生成器表达式中的嵌套 for 循环。 https://docs.python.org/3/library/itertools.html#itertools.product

import itertools

for i, j in itertools.product(range(3), range(4)):
    print(i, j)

Answer 3

一种方法是使用 np.mgrid to generate the indices and take advantage of numpy indexing (see here) 来避免任何 for 循环：

import numpy as np

array = np.random.rand(3, 4, 10)
x_vals = np.array([22, 150, 68, 41])  # size of 4 (same as size as 2nd axis of array above)

new_arr = np.zeros(array.shape)
rows, cols = map(np.ndarray.flatten, np.mgrid[0:3, 0:4])
new_arr[rows, :, cols] = array[rows, :, cols] * x_vals

Answer 4

只要打印是可选的，以下内容将以矢量化方式实现：

array = np.random.rand(3,4,10)

x_vals = np.array([22,150,68,41])

new_arr = array*x_vals.reshape((1,4,1))

这也假设您实际上想要 j 到运行一直到 10，而不是像您的示例那样停在 4

如何在 python 中重写这个嵌套的 for 循环？

How to rewrite this nested for loop in python?

python

numpy

cartesian-product

nested-loops

multidimensional-array