Numpy按列表切片疑问：如何提取行列对应的子矩阵

嘿，我完全懂你遇到的这个小坑——当你用非连续的行列索引列表提取子矩阵时，NumPy的默认花式索引会返回一维数组，而不是你想要的二维子矩阵。别担心，这里有两个简单靠谱的解决办法：

方法1：使用np.ix_()构建网格索引 #

np.ix_()函数专门用来把一维索引数组转换成能生成二维子矩阵的索引对，它会帮你把行和列的索引扩展成网格形式，让NumPy按笛卡尔积提取元素。

示例代码：

import numpy as np

# 初始化6×6单位矩阵
mat = np.diag((1,1,1,1,1,1))
# 目标行列索引
indices = [0,1,2,5]

# 提取4×4子矩阵
sub_matrix = mat[np.ix_(indices, indices)]
print(sub_matrix)

输出结果：

[[1 0 0 0]
 [0 1 0 0]
 [0 0 1 0]
 [0 0 0 1]]

方法2：分步提取行和列 #

另一种更直观的方式是先提取指定的行，再从结果中提取指定的列，这样也能得到目标子矩阵：

示例代码：

import numpy as np

mat = np.diag((1,1,1,1,1,1))
indices = [0,1,2,5]

# 先取指定行（得到4×6矩阵），再取指定列
sub_matrix = mat[indices][:, indices]
print(sub_matrix)

这个方法的输出和方法1完全一致，而且不需要额外调用np.ix_()，适合喜欢简洁写法的同学。

为什么原来的写法不行？ #

你之前用的mat[[0,1,2,5],[0,1,2,5]]属于配对式索引，NumPy会逐个取对应位置的元素：(0,0)、(1,1)、(2,2)、(5,5)，所以返回的是一维数组。而我们需要的是所有行索引和列索引的组合，也就是每一行都对应所有指定的列，上面两种方法正好实现了这个逻辑。

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