如何使用列表值作为索引对Numpy数组切片生成新的多维数组

解决方案 #

你只需要在调用np.take()时正确指定axis参数，再对列索引的结果做一次轴顺序调整即可，完整实现代码如下：

import numpy as np

A = np.array([[1,2,3],
              [4,5,6],
              [7,8,9]])
Cols, rows = np.array([[0,1],[1,2]])

# 行索引结果，直接指定axis=0取行维度即可，输出形状(2,2,3)符合预期
Row_results = np.take(A, rows, axis=0)

# 列索引结果，指定axis=1取列维度，再调整轴顺序得到形状(2,3,2)的预期结果
Cols_result = np.take(A, Cols, axis=1).transpose(1, 0, 2)

原理说明 #

• 行索引时，np.take沿着axis=0（行维度）索引，输入的rows形状为(2,2)，返回结果的维度顺序是「索引分组数、每组行数、每行列数」，正好和你需要的Row_results结构一致。
• 列索引时，np.take沿着axis=1（列维度）索引，默认返回的维度顺序是「行数、索引分组数、每组列数」，和你需要的结构顺序不一致，所以用transpose(1,0,2)把前两个维度交换位置，就能得到符合要求的Cols_result。

你也可以直接用Numpy高级索引实现相同效果：

# 高级索引实现列索引
Cols_result = A[:, Cols].transpose(1,0,2)
# 高级索引实现行索引
Row_results = A[rows, :]

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

内容的提问来源于stack exchange，提问作者E. hendrix