如何使用Numpy广播实现数组批量赋值以替代多行操作？

原代码中，state[:, pixels, :] = colors报错的核心原因是数组形状不兼容：

• state[:, pixels, :]的形状是(2, 3, 3)（2个batch、3个指定像素、每个像素3个颜色通道）
• colors的形状是(2, 3)，与目标区域的维度不匹配，无法直接广播。

要实现一行代码替代原两行赋值，只需将colors扩展一个维度，使其形状与目标区域兼容：

import numpy as np

state = np.arange(48, dtype='u1').reshape((2, 8, 3))
pixels = [3, 4, 5]
colors = [[42, 43, 44], [0, 1, 2]]

# 一行完成批量赋值
state[:, pixels, :] = np.array(colors)[..., np.newaxis]

原理说明 #

np.array(colors)[..., np.newaxis]将colors的形状从(2, 3)转换为(2, 3, 1)，此时Numpy会自动将其广播为(2, 3, 3)，与state[:, pixels, :]的形状完全匹配，实现：

• 第0个batch中，像素3/4/5的三个通道分别被42/43/44填充
• 第1个batch中，像素3/4/5的三个通道分别被0/1/2填充
  效果与原两行代码完全一致。

也可以用np.expand_dims(colors, axis=2)替代[..., np.newaxis]，实现相同的维度扩展。

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