如何使用索引列表对3D numpy数组执行正确切片操作

问题原因 #

你之前使用的arr[:,(1,3,4),(1,3,4)]属于Numpy的高级配对索引规则：当两个维度同时传入长度相同的序列作为索引时，Numpy会按位置一一匹配取值，最终得到的是(第一维长度, 索引序列长度)的二维数组，和你需要的提取指定行和列交叉区域的目标不符。 #

正确实现方案 #

两种方案均支持任意长度的索引列表，适配你索引数量大于3的实际使用场景：

方案1：分步索引（易读性最高） #

先截取第二维度的指定索引，再截取第三维度的指定索引，逻辑直观无需记忆特殊函数：

import numpy as np
np.random.seed(0)
arr = np.random.randint(0, 100, size=(2, 5, 5))
idx = [1,3,4]
# 先取第二维的idx行，再取第三维的idx列
k = arr[:, idx, :][:, :, idx]

方案2：使用np.ix_生成网格索引（性能最优） #

np.ix_可以将输入的索引序列转换为可广播的网格结构，一步完成交叉区域的提取，避免分步索引产生中间数组，性能更优：

idx = [1,3,4]
k = arr[:, np.ix_(idx, idx)]

两种方法输出的结果均和预期输出完全一致。 #

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