python - Numpy.where 与值列表一起使用

Question

我有一个二维和一维数组。我正在寻找包含至少一次来自 1d 数组的值的两行，如下所示:

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])

results = []

for elem in B:
    results.append(np.where(A==elem)[0])

这有效并产生以下数组:

[array([0, 5], dtype=int64),
 array([0, 3], dtype=int64),
 array([2, 4], dtype=int64),
 array([0, 2], dtype=int64)]

但这可能不是最好的处理方式。按照这个问题 ( Search Numpy array with multiple values ) 给出的答案，我尝试了以下解决方案:

out1 = np.where(np.in1d(A, B))

num_arr = np.sort(B)
idx = np.searchsorted(B, A)
idx[idx==len(num_arr)] = 0 
out2 = A[A == num_arr[idx]]

但是这些给了我不正确的值:

In [36]: out1
Out[36]: (array([ 0,  1,  2,  6,  8,  9, 13, 17], dtype=int64),)

In [37]: out2
Out[37]: array([0, 3, 1, 2, 3, 1, 2, 0])

谢谢你的帮助

Answer 1

如果您需要知道 A 的每一行是否包含数组 B 的任何元素而不关心它是 B 的哪个特定元素，可以使用以下脚本:

输入:

np.isin(A,B).sum(axis=1)>0 

输出:

array([ True, False,  True,  True,  True,  True])