如何使用 Numpy 向量化这个表达式？

Question

我正在尝试计算下面的数组“矩阵”。如何使用 NumPy 的向量化函数来做到这一点？

x = np.array([
    [2, 1, 0],
    [1, 1, 0],
    [3, 2, 1],
    [1, 0, 0],
    [2, 3, 0]
])

y = np.array([
    [405, 200, 150],
    [200, 300, 150],
    [150, 100, 105],
    [425, 200, 250],
    [500, 620, 300]
])

matrix = np.zeros((5,3,4))  
for i in range(5):
    for j in range(3):
       matrix[i,j,x[i,j]] = y[i,j]

我试过这个：

vmatrix = np.zeros((5,3,4))

vmatrix[:,:,x] = y

但是不起作用...

Answer 1

使用花式索引与网格：

ix, iy = np.mgrid[:5, :3]
matrix = np.zeros((5, 3, 4))
matrix[ix, iy, x] = y