Numpy数组提取单列后形状不符合预期的问题排查及解决方案

这是Numpy里新手很容易踩的一个小坑，我来给你理清楚问题出在哪，以及怎么解决～

首先要纠正一个小误解：你得到的[3 6]其实不是形状为(1,2)的二维数组，而是形状为(2,)的一维数组——Numpy在切片时会自动“挤压”掉长度为1的维度，这就是为什么z[:, 2]会返回一维数组，而不是你预期的(2,1)二维数组。

想要得到形状为(2,1)的列数组，同时满足后续和同形状变量相加的需求，有几种简单的方法：

方法1：用列表索引保留维度 #

当你用列表形式的索引（比如[:, [2]]）来提取列时，Numpy会保留原数组的二维结构，不会自动降维：

import numpy as np
z = np.array([[1,2,3], [4,5,6]])
target_col = z[:, [2]]
print(target_col)
# 输出：
# [[3]
#  [6]]
print(target_col.shape)  # 输出：(2, 1)

方法2：用reshape手动调整维度 #

如果你已经得到了一维数组，可以用reshape把它转换成目标形状，其中-1表示让Numpy自动计算该维度的长度，非常灵活：

target_col = z[:, 2].reshape(-1, 1)
print(target_col.shape)  # 输出：(2, 1)

方法3：用np.expand_dims扩展维度 #

这个函数可以明确指定要添加的维度位置，axis=1表示在列方向（第二个维度）添加一个长度为1的维度：

target_col = np.expand_dims(z[:, 2], axis=1)
print(target_col.shape)  # 输出：(2, 1)

完成上述操作后，你得到的target_col就可以和形状为(2,1)的变量正常执行相加操作啦，维度完全匹配，不会出现广播问题～

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