如何在NumPy中高效获取指定列全为1的行索引（无大量循环）

这问题用NumPy的原生向量化操作就能轻松解决，完全不需要写循环，效率还超高~

核心思路 #

我们需要两步完成筛选：

• 提取a1中由a2指定的列
• 检查这些列的每行是否全部等于1，再提取对应的行索引

完整代码实现 #

import numpy as np

a1 = np.array([[0, 0, 0], [1, 1, 0], [0, 1, 0], [1, 1, 1]])
a2 = np.array([1, 2])

# 1. 提取a2指定的列
selected_columns = a1[:, a2]
# 2. 创建掩码：标记每行是否所有指定列都为1
match_mask = (selected_columns == 1).all(axis=1)
# 3. 获取符合条件的行索引
matching_indices = np.where(match_mask)[0]

print(matching_indices)  # 输出结果：[3]

简化写法 #

如果想更简洁，可以把步骤合并成一行：

matching_indices = np.where((a1[:, a2] == 1).all(axis=1))[0]

为什么这个方法高效？ #

NumPy的向量化操作是基于C语言底层实现的，避免了Python循环的开销，哪怕处理超大数组也能保持很好的性能，这比手动写for循环遍历每一行要高效得多。

注意：如果你的a2是1-based的列编号（比如Excel里的列号），记得先转成0-based索引：a2 = a2 - 1，再代入上面的代码。

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