如何在NumPy中为数组添加指定大小的新维度而非默认大小1

核心原理说明 #

np.expand_dims、np.newaxis的作用仅为在数组形状中插入一个长度为1的维度，不会自动填充数据生成更长的维度，所以只能得到(2,3,1)这类结果。要生成指定长度的新维度，需要结合数据填充逻辑实现，以下是常用方案：

方案1：重复原数组填充新维度（所有新维度位置沿用原数组值） #

如果你的需求是新添加的第3维每个位置的值都和原二维数组对应位置的值相同，用np.repeat实现最直接：

import numpy as np
a = np.arange(6).reshape((2, 3))
# 先添加长度为1的第三维，再沿第三维重复5次
b = np.repeat(a[..., np.newaxis], repeats=5, axis=-1)
print(b.shape) # 输出 (2, 3, 5)

生成的数组示例：

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

       [[3, 3, 3, 3, 3],
        [4, 4, 4, 4, 4],
        [5, 5, 5, 5, 5]]])

方案2：初始化目标数组后赋值（新维度需要自定义不同值） #

如果新维度的每个切片需要填充不同的值，可以先创建指定形状的空数组，再利用numpy的广播特性把原二维数组的值赋值进去：

# 初始化形状为(2,3,5)的数组，dtype和原数组保持一致
c = np.zeros((a.shape[0], a.shape[1], 5), dtype=a.dtype)
# 广播赋值，原数组自动适配第三维的长度
c[:] = a[..., np.newaxis]
# 后续可以单独修改第三维任意切片的内容，比如 c[...,0] = 10

补充优化建议 #

如果后续只是要做数组间的数值运算，不需要显式生成(2,3,5)的数组：形状为(2,3,1)的数组可以直接和形状为(2,3,5)的数组做运算，numpy会自动广播适配，无需额外占用内存存储重复的数据。

内容的提问来源于stack exchange，提问作者Bilal Fawad Khan