如何查找2D NumPy数组中与1D数组匹配的行索引

问题原因 #

你之前的写法arr==a是逐元素进行广播对比，返回的是和arr形状相同的布尔数组，对应位置元素匹配则为True。直接传入np.where会返回所有匹配元素的坐标，第一个数组是匹配点的行索引，第二个是匹配点的列索引，因此会得到你看到的两组数组，而非整行匹配的结果。

正确实现方法 #

只需要在逐元素对比后，加按行全匹配判断即可，示例代码如下：

import numpy as np

arr = np.array([[1,2,3],
                [4,5,6],
                [7,8,9],
                [1,2,3]])
a = np.array([1,2,3])

# 先逐元素对比，再判断整行所有元素都匹配，最后提取行索引
match_rows = np.where((arr == a).all(axis=1))[0]
# 转成Python原生列表的话调用.tolist()
print(match_rows.tolist())
# 输出：[0, 3]

代码逻辑说明 #

• arr == a：利用NumPy广播特性，将一维数组a和arr的每一行逐元素对比，得到形状为(4,3)的布尔数组
• .all(axis=1)：指定沿列维度（即每行内部）做逻辑与运算，只有整行所有元素都匹配时才返回True，得到长度为4的一维布尔数组
• np.where(...)：提取布尔数组中True对应的索引，取第一个返回值就是匹配的行索引

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