Numpy where 与数组比较

Question

我有一个名为Y 的数组，其中包含类标签。我想找到与列表实验室指定的多个值匹配的 Y 的所有索引。

在这种情况下：

Y = np.array([1,2,3,1,2,3,1,2,3,1,2,3])
labs = [2,3]

我怎样才能做类似np.where(Y == labs) 的返回

array([1,2,4,5,7,8,10,11])

我知道一种可能性是遍历列表实验室并进行元素比较。但我正在寻找一种更加基于 pythonic/numpy 的解决方案，它可以避免循环。

Answer 1

您可以在 np.isin(..) [numpy-doc] 上使用 np.where(..) [numpy-doc]：

>>> np.where(np.isin(Y, L))[0]
array([ 1,  2,  4,  5,  7,  8, 10, 11])

.isin(Y, L) 将为我们提供True 和False 的数组，其中Y 的项目与L 中的元素匹配：

>>> np.isin(Y, labs)
array([False,  True,  True, False,  True,  True, False,  True,  True,
       False,  True,  True])

通过np.where(..)，我们将Trues 映射到相应的索引。

正如@hpaulj 所说，对于小的Ls，我们可以这样写：

np.any([Y == li for li in labs],axis=0)

在这里，对于labs 中的每个元素，我们将检查Y 是否是那个元素，并且我们使用np.any(..) 在它们之间创建一个“逻辑OR 链”以将其折叠为布尔值。