Question

我有一个ndarray A（例如2D图像），其值是从0到N的整数。

我还有另一个列表B或数组，其中包含从0到N范围内的数字的列表。

我想将第一个数组与第二个列表的每个元素进行比较，以获得一个新的ndarray，以指示pixel的值是否在列表中。

A约为10000 * 10000

B是一个具有10000-100000值的列表。

N达到50万

这是我希望获得的结果的一个例子。

我已经尝试过循环，它可以工作，但是由于矩阵很大，所以速度真的很慢。我还尝试使用.any()和numpy's比较功能来完成此操作，但未能获得所需的结果。

a = np.array([2, 23, 15, 0, 7, 5, 3])
b = np.array([3,7,17])
c = np.array([False, False, False, False, True, False, True])

Answer 1

您可以调整数组a的形状以使其具有额外的维度，该维度将用于与b进行比较，然后沿该维度使用np.any：

>>> np.any(a[..., None] == b, axis=-1)
array([False, False, False, False,  True, False,  True])

此方法非常灵活，因为它也可以与其他按元素比较的功能一起使用。例如，对于两个浮点数组，我们通常不希望np.equal来比较np.isclose，而可以通过简单地交换比较函数来做到这一点：

>>> np.any(np.isclose(a[..., None], b), axis=-1)

如果相等性是标准，那么np.isin的性能会更好，因为它不需要经过a.shape + (b.size,)形状的中间广播数组，该数组反而会沿最后一个轴减小。这意味着它既节省了内存，又节省了计算，因为它不需要分配该数组，也不需要执行所有计算：

In [2]: a = np.random.randint(0, 100, size=(100, 100))

In [3]: b = np.random.randint(0, 100, size=1000)

In [4]: %timeit np.any(a[..., None] == b, axis=-1)
12.1 ms ± 48.1 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

In [5]: %timeit np.isin(a, b)
608 µs ± 4.27 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

Answer 2

您可以使用numpy.in1d：

>>> np.in1d(a, b)
array([False, False, False, False,  True, False,  True], dtype=bool)

还有numpy.isin，推荐用于新代码。

如何快速比较ndarray的每个元素和排序列表/数组的每个元素？

2 个答案: