Numpy用布尔掩码选取指定行列生成新矩阵的更优实现方法

两种写法的结果差异原因 #

在Numpy的索引规则中，直接向两个维度分别传入一维布尔数组a[i, j]会触发逐元素配对索引逻辑：两个布尔数组按位置一一匹配，仅提取同时满足对应位置为True的坐标点元素，最终返回一维数组，也就是得到[1 8]的原因。
链式索引a[i, :][:, j]虽然能输出预期结果，但属于冗余写法，且中间步骤如果生成视图而非副本，后续操作可能引发非预期的修改问题，不推荐长期使用。

推荐的规范写法 #

方法1：使用np.ix_（首选通用方案） #

np.ix_是Numpy专门为多维度联合索引设计的工具，可将输入的一维索引序列转换为可广播的高维结构，写法简洁直观：

import numpy as np

a = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
i = np.array([True, False, True])
j = np.array([True, True, False])

# 一步得到目标二维矩阵
result = a[np.ix_(i, j)]
print(result)

输出结果：

[[1 2]
 [7 8]]

该写法支持任意数量维度的联合筛选，是Numpy官方推荐的标准实现。

方法2：手动升维触发广播 #

如果熟悉Numpy广播规则，也可以手动将行布尔数组扩展为列向量，和列布尔数组自动广播生成二维掩码实现筛选：

# 写法1：用np.newaxis扩展维度
result = a[i[:, np.newaxis], j]
# 写法2：用None简写扩展维度
result = a[i[:, None], j]

两种写法的运行效率和输出结果完全一致，可根据场景自行选择。

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