如何用NumPy向量化实现指定的三维数组赋值操作？

你的问题出在整数数组索引的维度匹配上，直接用vmatrix[:,:,x] = y会导致索引广播错误——因为x的形状是(5,3)，和前两个维度组合后会生成不匹配的索引维度，无法精准定位到目标位置。

正确的向量化实现方法 #

要实现和双重循环等价的向量化操作，需要生成与x、y形状匹配的i、j维度索引，通过三维整数数组索引精准定位赋值位置：

import numpy as np

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

vmatrix = np.zeros((5,3,4))
# 生成i和j的网格索引，广播后与x、y形状一致
i_indices = np.arange(5)[:, np.newaxis]  # 形状(5,1)，广播为(5,3)
j_indices = np.arange(3)                 # 形状(3,)，广播为(5,3)
# 精准赋值：每个(i,j,x[i,j])位置对应y[i,j]
vmatrix[i_indices, j_indices, x] = y

原理说明 #

• i_indices[:, np.newaxis]将一维数组[0,1,2,3,4]转换为列向量，和j_indices广播后，两者形状都变为(5,3)，完全匹配x和y的维度。
• 此时vmatrix[i_indices, j_indices, x]会直接定位到每个(i,j)对应的第三维索引x[i,j]，赋值逻辑和原循环完全等价，且是NumPy原生向量化操作，效率远高于Python循环。

验证结果 #

你可以通过以下代码确认结果和原循环一致：

# 原循环生成的matrix
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]

# 检查是否相等
print(np.array_equal(matrix, vmatrix))  # 输出True

内容的提问来源于stack exchange，提问作者MBlrd