如何优化使用numpy.where查找含一维数组元素的二维数组行？

嘿，你的思路是对的——循环遍历B的每个元素再逐个查找确实不是最优解，Numpy的向量化操作能帮我们大幅提升效率，尤其是当数组规模较大的时候。

问题分析 #

你的需求是找出二维数组A中至少包含B中任意一个元素的行。原代码通过循环每个B元素并调用np.where，得到的是每个元素对应的行索引列表，但不仅有重复计算，还需要后续手动去重才能得到最终的目标行索引。我们可以一步到位完成这个任务。

优化后的代码 #

import numpy as np
A = np.array([[0, 3, 1], [9, 4, 6], [2, 7, 3], [1, 8, 9], [6, 2, 7], [4, 8, 0]])
B = np.array([0,1,2,3])

# 生成布尔矩阵：每个元素标记是否属于B
element_in_B = np.isin(A, B)
# 按行判断：只要该行有一个元素属于B，就标记为True
row_mask = element_in_B.any(axis=1)
# 获取所有满足条件的行索引
target_rows = np.where(row_mask)[0]

print(target_rows)  # 输出: array([0, 2, 3, 4, 5])

为什么这更高效？ #

• 避免Python循环：Numpy的底层是C实现的向量化操作，比Python级别的循环快几个数量级，尤其是当A和B的规模变大时，差距会非常明显。
• 一步到位：直接生成行级别的判断掩码，不需要收集多个索引数组再合并去重，代码更简洁，逻辑更清晰。

如果需要保留原代码的输出格式（每个B元素对应的行索引） #

如果你确实需要得到原代码那样的“每个B元素对应的行索引列表”，也可以用向量化的方式减少循环开销：

# 广播A和B，生成形状为(len(B), *A.shape)的布尔数组
matches = (A == B[:, np.newaxis, np.newaxis])
# 按行聚合，得到每个B元素对应的行索引
result = [np.where(matches[i].any(axis=1))[0] for i in range(len(B))]
print(result)
# 输出: [array([0, 5]), array([0, 3]), array([2, 4]), array([0, 2])]

不过这个方法仅在你需要逐个元素的结果时使用，如果你只需要所有满足条件的行索引，第一种方法是绝对最优选择。

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