Numpy中np.where匹配多值返回结果长度为1的问题求解

你猜的完全正确！np.where() 返回的结果是一个元组，这是它的设计特性——为了兼容多维数组的索引返回需求。比如处理二维数组时，它会返回包含行索引数组和列索引数组的元组，格式像(row_indices, col_indices)。而你这里是一维数组，所以元组里只有一个元素，就是你要的索引数组。

两种快速提取目标数组的方法： #

方法1：直接提取元组的第一个元素

只需要在np.where()的结果后面加上[0]，就能拿到里面的索引数组：

import numpy as np
a = [1,3,5,7,9,11,13,15]
b = [1,7,13]
x = np.where(np.isin(a,b))[0]
print(x)       # 输出：[0 3 6]
print(len(x))  # 输出：3

方法2：使用np.flatnonzero()简化代码

如果你只需要一维数组的非零元素索引（这里np.isin(a,b)返回的布尔数组里，匹配的位置是True，等价于1），可以直接用np.flatnonzero()，它会直接返回一维的索引数组，省去拆元组的步骤：

x = np.flatnonzero(np.isin(a,b))
print(x)       # 输出：[0 3 6]

额外说明 #

如果以后处理多维数组，比如二维数组：

a_2d = np.array([[1,2],[3,4],[5,6]])
b = [2,5]
x = np.where(np.isin(a_2d,b))
print(x)  # 输出：(array([0, 2]), array([1, 0]))

这时候元组里的两个数组分别对应匹配元素的行索引和列索引，你可以根据需求单独提取其中一个，比如x[0]拿行索引，x[1]拿列索引。

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