NumPy：比较两个数组中的元素

Question

有没有人遇到这个问题？假设你有两个如下的数组：

a = array([1,2,3,4,5,6])
b = array([1,4,5])

有没有办法比较a中存在的元素是否存在于b中？例如，

c = a == b # Wishful example here
print c
array([1,4,5])
# Or even better
array([True, False, False, True, True, False])

我正在尝试避免使用循环，因为如果有数百万个元素的话，这将需要很长时间。 有什么建议吗？

谢谢

- ebressert

---

你的数组中有什么数据？类似于示例中的唯一整数索引吗？ - u0b34a0f6ae

Answer 1

实际上，比这些解决方案更简单的解决方案是：

import numpy as np

a = array([1,2,3,4,5,6])
b = array([1,4,5])

c = np.in1d(a,b)

得到的 c 是:

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

- eteq

---

有没有一个“几乎相等”的版本？你可以指定用于测试相等性的条件吗？ - endolith

现在这个已经被废弃了。NumPy文档指出：“我们建议对于新代码使用isin代替in1d”。 - divenex

Answer 2

使用 np.intersect1d。

#!/usr/bin/env python
import numpy as np
a = np.array([1,2,3,4,5,6])
b = np.array([1,4,5])
c=np.intersect1d(a,b)
print(c)
# [1 4 5]

请注意，如果a或b具有非唯一元素，np.intersect1d将给出错误的答案。在这种情况下，请使用np.intersect1d_nu。

还有np.setdiff1d、setxor1d、setmember1d和union1d。请参见 Numpy Example List With Doc。

- unutbu

Answer 3

感谢您的回复，kaizer.se。虽然不完全符合我的要求，但是通过朋友的建议和您的意见，我想出了以下方案。

import numpy as np

a = np.array([1,4,5]).astype(np.float32)
b = np.arange(10).astype(np.float32)

# Assigning matching values from a in b as np.nan
b[b.searchsorted(a)] = np.nan

# Now generating Boolean arrays
match = np.isnan(b)
nonmatch = match == False

这个过程有点繁琐，但比使用循环或在循环中使用Weave要好。

干杯！

- ebressert

---

这种方法的问题在于它即使对于a中不存在于b中的值也会返回索引（以及其他情况下的重复索引）。例如：numpy.searchsorted([1, 2], [1.2, 1.3])返回[1, 1]，这对于OP来说是不合适的。 - sirfz

Answer 4

Numpy有一个函数numpy.setmember1d()，可以用于已排序且去重的数组，返回你所需的布尔数组。如果输入数组不符合条件，则需要将其转换为set格式并对结果进行反转换。

import numpy as np
a = np.array([6,1,2,3,4,5,6])
b = np.array([1,4,5])

# convert to the uniqued form
a_set, a_inv = np.unique1d(a, return_inverse=True)
b_set = np.unique1d(b)
# calculate matching elements
matches = np.setmea_set, b_set)
# invert the transformation
result = matches[a_inv]
print(result)
# [False  True False False  True  True False]

编辑：很遗憾，numpy中的setmember1d方法非常低效。您提议的搜索排序和分配方法速度更快，但是如果可以直接分配结果，最好直接分配结果并避免大量不必要的复制。此外，如果b包含a中没有的任何内容，则您的方法将失败。以下纠正了这些错误：

result = np.zeros(a.shape, dtype=np.bool)
idxs = a.searchsorted(b)
idxs = idxs[np.where(idxs < a.shape[0])] # Filter out out of range values
idxs = idxs[np.where(a[idxs] == b)] # Filter out where there isn't an actual match
result[idxs] = True
print(result)

我的基准测试显示，使用您的方法需要6.6毫秒，而使用numpy setmember1d在1M个元素a和100个元素b上需要109毫秒，而我的方法只需要91微秒。

- Ants Aasma

---

这是一个不错的解决方案。我会尝试你的建议和我刚刚写的内容，看看哪个速度更优。非常感谢大家的帮助！ - ebressert

我写的方法速度稍快。对于一个10000元素的数组，在iPython中使用timeit所需的时间大约为3微秒。而setmember1d方法需要3毫秒。我认为你的方法更优雅，但我需要速度。 - ebressert

你在第三行忘记了关闭一个括号，你应该在某位计算机科学教授注意到之前修复它... - dalloliogm

ebressert：看起来你是对的，numpy中的setmember1d实现非常糟糕。但是你使用的方法似乎没有充分的理由使用nan值，你可以直接使用结果数组。我将编辑并提供相应的示例。 - Ants Aasma

Ants Aasma：你的修改很好。我将其中的一部分实现到我的代码中，又一次提高了速度。我没有使用nans，而是使用了-1，并在匹配=b>=0时进行了过滤。在我的情况下，我正在处理索引，因此没有-1的索引。这就是为什么我使用np.nan，它适用于更一般的情况。感谢您的建议。我的代码现在真的很快。 - ebressert

Answer 5

ebresset，你的答案 只有在 a 是 b 的子集（且 a 和 b 已排序）时才能起作用。否则，searchsorted 将返回错误的索引。我曾经做过类似的事情，并将其与你的代码结合使用：

# Assume a and b are sorted
idxs = numpy.mod(b.searchsorted(a),len(b))
idxs = idxs[b[idxs]==a]
b[idxs] = numpy.nan
match = numpy.isnan(b)

- AFoglia

Answer 6

你的例子暗示了集合的行为，更关心在数组中的存在性，而不是正确的元素放在正确的位置。Numpy使用其数学数组和矩阵以不同的方式处理此问题，它只会告诉你确切位置上的项目。你能让它对你起作用吗？

>>> import numpy
>>> a = numpy.array([1,2,3])
>>> b = numpy.array([1,3,3])
>>> a == b
array([ True, False,  True], dtype=bool)

- u0b34a0f6ae

---

抱歉，如果您尝试此示例，它将无法正常工作；此外，您还需要先对数组进行排序。 - dalloliogm

@dalloligom：嗯，我是从我的交互式会话中复制的，所以至少对于某些版本的Python和Numpy来说，它的工作方式确实完全相同。 - u0b34a0f6ae

好的，但是如果两个数组长度不同，它就无法工作；无论如何，您都需要首先对它们进行排序（尝试array([1,2,3])==array([2,3,1])）。他想知道一个数组中的哪些元素存在于另一个数组中。 - dalloliogm

顺便说一下，即使对数组进行排序也行不通...你必须使用一个集合结构。 - dalloliogm

@dalloliogm：你看了我的回答吗？我似乎没有理解所有的内容吗？ - u0b34a0f6ae