Numpy 4D数组索引错误：如何选取指定列规避单省略号限制

我有如下NumPy 4D数组：

import numpy as np
X = np.random.rand(5, 1, 10, 4)

# 前2个元素示例：
X[:2]
# 输出：
array([[[[0.27383924, 0.48908027, 0.64997038, 0.20394247],
         [0.28361942, 0.33425344, 0.27687327, 0.2549442 ],
         [0.91655337, 0.77325791, 0.31945728, 0.82919328],
         [0.83989813, 0.65384396, 0.13853182, 0.46299719],
         [0.14700217, 0.37591964, 0.8545056 , 0.02064633],
         [0.06186759, 0.88515897, 0.84535195, 0.41697788],
         [0.9180413 , 0.42174186, 0.55005076, 0.70799608],
         [0.68446734, 0.41968608, 0.19013073, 0.16875907],
         [0.44687274, 0.62684239, 0.27798323, 0.6355134 ],
         [0.8489883 , 0.23450424, 0.53215137, 0.66814813]]],


       [[[0.85473496, 0.70600538, 0.70862705, 0.89358703],
         [0.80026841, 0.62795239, 0.06190375, 0.41356739],
         [0.01792312, 0.82979946, 0.82117873, 0.14904196],
         [0.10831188, 0.63943446, 0.20393167, 0.4058673 ],
         [0.7966648 , 0.37533761, 0.73456441, 0.36841977],
         [0.78459342, 0.34400906, 0.08502799, 0.2625697 ],
         [0.57079306, 0.52439791, 0.6417777 , 0.02517128],
         [0.84525549, 0.40980805, 0.20189425, 0.39604223],
         [0.06425004, 0.75075354, 0.69504595, 0.76566498],
         [0.01929747, 0.03261916, 0.32740129, 0.43836062]]]]])

选取每个元素的前两列可以用X[..., :2]，示例如下：

X[..., :2][:2]
# 输出：
array([[[[0.27383924, 0.48908027],
         [0.28361942, 0.33425344],
         [0.91655337, 0.77325791],
         [0.83989813, 0.65384396],
         [0.14700217, 0.37591964],
         [0.06186759, 0.88515897],
         [0.9180413 , 0.42174186],
         [0.68446734, 0.41968608],
         [0.44687274, 0.62684239],
         [0.8489883 , 0.23450424]]],


       [[[0.85473496, 0.70600538],
         [0.80026841, 0.62795239],
         [0.01792312, 0.82979946],
         [0.10831188, 0.63943446],
         [0.7966648 , 0.37533761],
         [0.78459342, 0.34400906],
         [0.57079306, 0.52439791],
         [0.84525549, 0.40980805],
         [0.06425004, 0.75075354],
         [0.01929747, 0.03261916]]]]])

现在需要选取前两列和最后一列（剔除第三列），执行X[..., :2, ...,3]会触发错误：IndexError: an index can only have a single ellipsis ('...')，期望输出如下：

# X的前2个元素的情况
array([[[[0.27383924, 0.48908027,  0.20394247],
         [0.28361942, 0.33425344,  0.2549442 ],
         [0.91655337, 0.77325791,  0.82919328],
         [0.83989813, 0.65384396,  0.46299719],
         [0.14700217, 0.37591964,  0.02064633],
         [0.06186759, 0.88515897,  0.41697788],
         [0.9180413 , 0.42174186,  0.70799608],
         [0.68446734, 0.41968608,  0.16875907],
         [0.44687274, 0.62684239,  0.6355134 ],
         [0.8489883 , 0.23450424,  0.66814813]]],


       [[[0.85473496, 0.70600538,  0.89358703],
         [0.80026841, 0.62795239,  0.41356739],
         [0.01792312, 0.82979946,  0.14904196],
         [0.10831188, 0.63943446,  0.4058673 ],
         [0.7966648 , 0.37533761,  0.36841977],
         [0.78459342, 0.34400906,  0.2625697 ],
         [0.57079306, 0.52439791,  0.02517128],
         [0.84525549, 0.40980805,  0.39604223],
         [0.06425004, 0.75075354,  0.76566498],
         [0.01929747, 0.03261916,  0.43836062]]]]])

解决方案 #

方法1：直接指定列索引列表

数组最后一维的列索引为0、1、2、3，我们需要保留0、1、3列，直接用索引列表选取即可：

result = X[..., [0, 1, 3]]
# 查看前2个元素验证
result[:2]

该方法会直接返回包含指定列的数组，完全符合期望输出。

方法2：拼接前两列与最后一列

先分别提取前两列和最后一列，再通过np.concatenate在最后一维拼接：

# 提取前两列
first_two_cols = X[..., :2]
# 提取最后一列（用-1:保持维度一致，避免降维）
last_col = X[..., -1:]
# 拼接数组
result = np.concatenate([first_two_cols, last_col], axis=-1)
# 查看前2个元素验证
result[:2]

这种方法适合需要分步处理的场景，同样能得到目标结果。

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