Python3：向量化嵌套循环答案

【问题标题】：Python3: vectorizing nested loopsPython3：向量化嵌套循环
【发布时间】：2020-04-03 19:51:39
【问题描述】：

我有这个功能：

def fun(x):   # x is a vector with size: (size_x*size_y) = n 
    c = 0
    f_vec = np.zeros((size_x*size_y))


    for i in range(size_x):
        for j in range(size_y):
             f_vec[c]=i*j*x[c]   
             c=c+1

    return f_vec

我这样做是因为向量 x 是（考虑 size_x=4 和 size_y=3）

 x[0]=x00    #c=0  i=0,j=0
 x[1]=x01    #c=1  i=0, j=1
 x[2]=x02    #c=2   i=0. j=size_y-1
 x[3]=x10    #c=3   i=1, j=0
 x[4]=x11
  ...
 x[n]=x32    #c=n   i=size_x-1, j= size_y-1

我可以避免嵌套循环并进行简单的向量运算吗？我想要类似 f[c] = F[x[c]] *i *j

但是通过知道 c 的值来找到 i 和 j 并不是那么简单。你知道方法吗？

谢谢。

【问题讨论】：

标签： python-3.x numpy vectorization

【解决方案1】：

您可以为此使用广播：

(
    x.reshape(size_x, size_y) *
    np.arange(size_x)[:, None] *
    np.arange(size_y)
).ravel()

或爱因斯坦求和形式

np.einsum(
    'ij,i,j->ij',
    x.reshape(size_x, size_y),
    np.arange(size_x),
    np.arange(size_y)
).ravel()

【讨论】：

这不适用于更复杂的功能，对吗？假设您不是简单地返回 i*j，而是有依赖于 i 和 j 的条件和其他变量。我可以访问每个计算的 i 和 j 的值吗？
Depends ;-) 只需使用代码 sn-p 打开另一个问题，显示您正在尝试做什么
我创建了一个新问题。这是链接：link

【解决方案2】：

本质上，这与Nils Werner's answer相同，但我发现将i*j部分理解为2D ndarray np.outer(np.arange(x_size), np.arange(y_size)更容易，然后进行广播：

(x.reshape(x_size, y_size) * np.outer(np.arange(x_size), np.arange(y_size)).ravel()

如果您针对x_size 和y_size 的相同值重复执行此操作，则可以预先计算i*j 部分。

【讨论】：