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

Question

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

我尝试使用zip python 中的函数，但它不打印出 i,j 的顺序与嵌套循环相同的方式，这是我的计算所必需的。

这个问题与 stackoverflow 的问题类似，但我很难重现答案: 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，相当于生成器表达式中的嵌套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)