在 NumPy 数组中插入一个新轴

在 NumPy 中，我们经常需要在数组的某个位置插入一个新的轴。插入新轴的机制是通过 NumPy 函数 np.newaxis 来实现。

插入一个新轴

我们可以使用 np.newaxis 来在 NumPy 数组中插入一个新轴。新轴的位置可以在任何维度上插入，下面是一个将一维数组转化为二维数组的例子：

import numpy as np

x = np.array([1, 2, 3, 4, 5])
print("原数组 x 的形状：", x.shape)

# 在第一维上插入新轴，将一维数组转化为二维数组，x.T 是原数组的转置
y = x[np.newaxis, :]
print("插入新轴之后数组 y 的形状：", y.shape)
print("数组 y 的值：\n", y)

# 在第二维上插入新轴
z = x[:, np.newaxis]
print("插入新轴之后数组 z 的形状：", z.shape)
print("数组 z 的值：\n", z)

输出：

原数组 x 的形状： (5,)
插入新轴之后数组 y 的形状： (1, 5)
数组 y 的值：
 [[1 2 3 4 5]]
插入新轴之后数组 z 的形状： (5, 1)
数组 z 的值：
 [[1]
  [2]
  [3]
  [4]
  [5]]

删除一个轴

使用 np.squeeze() 函数可以压缩数组中长度为 $1$ 的维度。例如：

a = np.array([[[0], [1], [2]]])
print("原数组 a 的形状：", a.shape)

# 删除长度为 1 的一维，a.squeeze() 可以省略参数，默认删除所有长度为 1 的维度
b = np.squeeze(a)
print("压缩维度之后数组 b 的形状：", b.shape)
print("数组 b 的值：\n", b)

输出：

原数组 a 的形状： (1, 3, 1)
压缩维度之后数组 b 的形状： (3,)
数组 b 的值：
 [0 1 2]

小结

通过 np.newaxis，我们可以轻松地在 NumPy 数组中插入新维度。如果在数组中存在长度为 $1$ 的维度，可以使用 np.squeeze() 函数来删除这些维度，从而使数组更紧凑。