Python3:向量化嵌套循环 [英] Python3: vectorizing nested loops

问题描述

我有这个功能:

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)

I do this because what happens is that the vector x is (considering size_x=4 and 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

Can I avoid the nested loop and do a simple vector operation? I would like to have something like f[c] = F[x[c]] *i *j

但是通过知道c值来找到i和j并不是那么简单. 你知道吗?

But it is not that simple to find i and j by knowing the c value. Do you know a way?

谢谢.

推荐答案

您可以为此使用广播:

(
    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()

这篇关于Python3:向量化嵌套循环的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！