使用numpy进行元素"in"的Pythonic和有效方法

Question

我正在寻找一种方法来有效地获得一个布尔数组,其中给定两个具有相同大小的数组,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 的广播规则，您应该首先对数组进行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)您将获得一个二维布尔数组，显示每个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)