Numpy多维数组索引异常：如何保留维度与轴序进行子选择

问题描述 #

我有一个多维数组，以及两个整数列表L_i和L_j，分别对应需要保留的第i轴和第j轴元素。要求：

• 保留数组原维度，即使L_i或L_j仅包含一个元素（不折叠单例轴）；
• 保留轴的原有顺序。

以下是复现意外行为的示例：

import numpy as np
aa = np.arange(120).reshape(5,4,3,2)
aa.shape
# 预期结果为(5,4,3,2)

aa[:,:,:,[0,1]].shape
# 预期结果为(5, 4, 3, 2)

aa[:,:,:,[0]].shape
# 得到(5,4,3,1)符合预期，即使仅选一个元素仍保留最后一维

aa[:,[1,3],:,[0]].shape
# 得到(2, 5, 3)，与预期的(5, 2, 3, 1)不符！

想了解Numpy为何会折叠并重新排序轴，以及正确进行子选择的最佳方法。

原因分析 #

问题源于Numpy的高级索引特性：

• 当索引中同时存在多个一维数组（如示例中的[1,3]和[0]）时，Numpy会触发高级索引模式。此时，所有使用数组索引的轴会被前置，切片（:）对应的轴被后置，再结合广播机制，最终导致轴顺序混乱、维度折叠。
• 单个一维数组索引（如[0]）属于基础索引+新轴保留，仅在对应轴选择元素并保留维度，不会改变轴顺序。

最佳解决方案 #

方法1：使用np.take（推荐） #

np.take可针对指定轴选取元素，通过keepdims=True强制保留原维度，且完全维持轴顺序：

import numpy as np

aa = np.arange(120).reshape(5,4,3,2)
# 示例：选择第1轴的[1,3]元素、第3轴的[0]元素
result = np.take(np.take(aa, [1,3], axis=1, keepdims=True), [0], axis=3, keepdims=True)
print(result.shape)  # 输出 (5, 2, 3, 1)，符合预期

• 优势：逻辑直观，直接针对目标轴操作，无轴顺序混乱风险。

方法2：结合np.expand_dims与轴交换 #

若偏好方括号索引，可先获取结果，再补全维度并调整轴顺序：

temp = aa[:,[1,3],:,[0]]
# 交换轴0和轴1，再在末尾添加单例轴
result = np.expand_dims(temp.swapaxes(0,1), axis=-1)
print(result.shape)  # 输出 (5, 2, 3, 1)

• 劣势：需手动调整轴顺序，适合简单场景。

方法3：使用np.ix_构建索引 #

np.ix_可将一维数组转换为广播兼容的索引网格，确保轴顺序不变：

# 构建索引列表，slice(None)对应原轴的全选
idx = [slice(None), [1,3], slice(None), [0]]
# 用np.ix_处理数组索引部分，维持轴维度与顺序
idx[1] = np.ix_(idx[1])[0]
idx[3] = np.ix_(idx[3])[0]
result = aa[tuple(idx)]
print(result.shape)  # 输出 (5, 2, 3, 1)

• 优势：可直接基于原索引逻辑修改，适配复杂多轴选择场景。

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