时间戳

首页 > 解决方案 > 在 numpy 中向量化基于索引的矩阵运算

python - 在 numpy 中向量化基于索引的矩阵运算

问题描述

如何矢量化以下循环？

def my_fnc():
   m = np.arange(27.).reshape((3,3,3))
   ret = np.empty_like(m)

   it = np.nditer(m, flags=['multi_index'])
   for x in it:
      i,j,k = it.multi_index
      ret[i,j,k] = x / m[i,j,i]
   return ret

基本上，我将 m 中的每个值除以类似于对角线的东西。并非 m 中的所有值都会不同，arange 只是一个示例。

提前致谢！~

PS：这是上面函数的输出，不要介意nans :)

array([[[        nan,         inf,         inf],
        [ 1.        ,  1.33333333,  1.66666667],
        [ 1.        ,  1.16666667,  1.33333333]],

       [[ 0.9       ,  1.        ,  1.1       ],
        [ 0.92307692,  1.        ,  1.07692308],
        [ 0.9375    ,  1.        ,  1.0625    ]],

       [[ 0.9       ,  0.95      ,  1.        ],
        [ 0.91304348,  0.95652174,  1.        ],
        [ 0.92307692,  0.96153846,  1.        ]]])

标签： pythonnumpyvectorization

解决方案

用于一次性advanced-indexing获得m[i,j,i]等价物，然后简单地将输入数组除以它 -

r = np.arange(len(m))
ret = m/m[r,:,r,None] # Add new axis with None to allow for broadcasting

推荐阅读