使用 numpy 执行 elementwise "in"的 Pythonic 和有效方法

Question

我正在寻找一种有效获取 boolean 值数组的方法，其中给定两个大小相等的数组 a 和 b，如果对应的每个元素都为真a 的元素出现在 in b 的相应元素中。

例如下面的程序:

a = numpy.array([1, 2, 3, 4])
b = numpy.array([[1, 2, 13], [2, 8, 9], [5, 6], [7]])
print(numpy.magic_function(a, b))

应该打印

[True, True, False, False]

记住这个函数应该等同于

[x in y for x, y in zip(a, b)]

只有 numpy - 针对 a 和 b 很大的情况进行了优化，并且 b 的每个元素都是相当小。

Answer 1

利用 NumPy 的 broadcasting规则你应该首先使数组 b 平方，这可以使用 itertools.izip_longest 来实现:

from itertools import izip_longest

c = np.array(list(izip_longest(*b))).astype(float)

导致:

array([[  1.,   2.,   5.,   7.],
       [  2.,   8.,   6.,  nan],
       [ 13.,   9.,  nan,  nan]])

然后，通过执行 np.isclose(c, a)，您将获得一个二维 boolean 数组，显示每个 c[:, i] 和 之间的差异>a[i]，根据广播规则，给出:

array([[ True,  True, False, False],
       [False, False, False, False],
       [False, False, False, False]], dtype=bool)

可用于获取您的答案:

np.any(np.isclose(c, a), axis=0)
#array([ True,  True, False, False], dtype=bool)