如何将从4D numpy数组索引得到的2D数组转为3D数组？

问题分析 #

你现在碰到的情况是：np.where返回的是两个一维索引数组，当用它们去索引dataX的后两个维度时，NumPy会把这两组索引一一对应，把所有匹配到的元素“扁平化”成一个一维序列，所以最终得到的element是(4, 860498)的2D数组。要转换成3D数组，核心就是给提取后的结果增加一个额外维度，或者调整索引方式来保留维度信息。

解决方案 #

这里有几种简单易实现的方法：

方法1：索引后直接添加维度

用np.newaxis（和None等价）给结果加一个维度，比如把新维度放在最后，得到(4, 860498, 1)的3D数组：

import numpy as np

# 你的原有代码
idx = np.where(dataY == 5.0)
# 添加新维度到末尾
element = dataX[:, idx[0], idx[1]][:, :, np.newaxis]

如果想把新维度放在中间（形状变为(4, 1, 860498)），只需要调整np.newaxis的位置：

element = dataX[:, idx[0], idx[1]][:, np.newaxis, :]

方法2：用np.expand_dims显式扩展维度

np.expand_dims是更直观的扩展维度工具，你可以明确指定要添加的轴的位置：

# 在最后添加维度，得到(4, 860498, 1)
element_3d = np.expand_dims(dataX[:, idx[0], idx[1]], axis=2)
# 在中间添加维度，得到(4, 1, 860498)
element_3d = np.expand_dims(dataX[:, idx[0], idx[1]], axis=1)

方法3：调整索引的维度结构

在索引前就把一维索引转换成二维结构，这样提取时会自然生成3D数组：

# 给每个索引数组添加一个新维度，变成(860498, 1)的形状
idx_rows = idx[0][:, np.newaxis]
idx_cols = idx[1][:, np.newaxis]
# 索引后直接得到(4, 860498, 1)的3D数组
element = dataX[:, idx_rows, idx_cols]

验证结果 #

不管用哪种方法，你都可以通过print(element.shape)来确认最终数组的形状是否符合预期。

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