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

Question

如何矢量化以下循环？

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 中的所有值都会不同，范围只是一个示例。

提前致谢! ~

P.S.:这是上面函数的输出，不要介意 nan :)

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.        ]]])

Answer 1

使用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