如何基于NumPy数组第一维度为最后维度元素赋值？

要实现你想要的赋值效果，核心是利用NumPy的广播机制，让输入的一维数组[1,2,3]匹配目标区域的形状(3,2,2)。

单行实现代码 #

import numpy as np

a = np.array([
  [[[0, 0], [0, 0]], 
    [[0, 0], [0, 0]]],
  [[[0, 0], [0, 0]], 
    [[0, 0], [0, 0]]],
  [[[0, 0], [0, 0]], 
    [[0, 0], [0, 0]]],
])

# 方法1：用reshape扩展维度
a[:, ..., 1] = np.array([1, 2, 3]).reshape(3, 1, 1)

# 方法2：更简洁的维度扩展写法
# a[:, ..., 1] = np.array([1, 2, 3])[:, None, None]

print(a)

运行结果 #

array([[[[0, 1],
         [0, 1]],

        [[0, 1],
         [0, 1]]],


       [[[0, 2],
         [0, 2]],

        [[0, 2],
         [0, 2]]],


       [[[0, 3],
         [0, 3]],

        [[0, 3],
         [0, 3]]]])

原理说明 #

你之前的代码报错是因为：目标赋值区域a[:, ..., 1]的形状是(3,2,2)，而输入的[1,2,3]形状是(3,)，两者维度不匹配，无法直接广播。

通过把[1,2,3]扩展为(3,1,1)的形状（两种方法都能实现），NumPy会自动将每个元素广播到后面的2x2维度，正好对应每个第一维度下所有子数组的最后一维第二个元素，完美匹配你的需求。

对比错误尝试 #

你用[1,2]赋值时不报错，是因为(2,)的数组会被广播到(3,2,2)——它会沿着第一维度重复3次，所以每个第一维度的子数组都用[1,2]填充，这显然不是你想要的效果。

这种向量化操作比Python循环高效得多，完全符合你要的简洁优雅的单行实现需求。

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