在另一个数组中查找一个数组的匹配索引

问题描述：

我有两个 numpy 数组，A 和 B。A 包含唯一值，B 是 A 的子数组。现在我正在寻找一种方法来获取 A 中 B 值的索引。

例如：

A = np.array([1,2,3,4,5,6,7,8,9,10])
B = np.array([1,7,10])
# I need a function fun() that:
fun(A,B)
>> 0,6,9

解决方案 1：

您可以np.in1d使用np.nonzero-

np.nonzero(np.in1d(A,B))[0]

np.searchsorted如果您关心维护顺序，也可以使用-

np.searchsorted(A,B)

对于一般情况，当A&B是未排序数组时，您可以引入sorter中的选项np.searchsorted，如下所示 -

sort_idx = A.argsort()
out = sort_idx[np.searchsorted(A,B,sorter = sort_idx)]

我也会添加我最喜欢的broadcasting来解决一般情况 -

np.nonzero(B[:,None] == A)[1]

样本运行 -

In [125]: A
Out[125]: array([ 7,  5,  1,  6, 10,  9,  8])

In [126]: B
Out[126]: array([ 1, 10,  7])

In [127]: sort_idx = A.argsort()

In [128]: sort_idx[np.searchsorted(A,B,sorter = sort_idx)]
Out[128]: array([2, 4, 0])

In [129]: np.nonzero(B[:,None] == A)[1]
Out[129]: array([2, 4, 0])

解决方案 2：

你试过嗎searchsorted？

A = np.array([1,2,3,4,5,6,7,8,9,10])
B = np.array([1,7,10])

A.searchsorted(B)
# array([0, 6, 9])

解决方案 3：

只是为了完整性：如果中的值A非负且相当小：

lookup = np.empty((np.max(A) + 1), dtype=int)
lookup[A] = np.arange(len(A))
indices  = lookup[B]

解决方案 4：

我最近也有同样的问题。但是，时间性能对我来说非常关键。因此，我猜不同解决方案的时间比较可能对其他人有用。

正如 Divakar 所提到的，您可以使用np.in1d(A, B)和np.where。np.nonzero此外，您还可以使用np.in1d(A, B)和np.intersect1d（基于此页面）。此外，您还可以将其用作np.searchsorted对数组进行排序的另一种有用方法。

我想添加另一个简单的解决方案。您可以使用理解列表。它可能比以前的解决方案花费的时间更长。但是，如果您利用Numba python 包，则耗时会少得多。

In [1]: import numpy as np
In [2]: from numba import njit

In [3]: a = np.array([1,2,3,4,5,6,7,8,9,10])
In [4]: b = np.array([1,7,10])

In [5]: np.where(np.in1d(a, b))[0]
   ...: array([0, 6, 9])

In [6]: np.nonzero(np.in1d(a, b))[0]
   ...: array([0, 6, 9])   

In [7]: np.searchsorted(a, b)
   ...: array([0, 6, 9])


In [8]: np.searchsorted(a, np.intersect1d(a, b))
   ...: array([0, 6, 9])

In [9]: [i for i, x in enumerate(a) if x in b]
   ...: [0, 6, 9]

In [10]: @njit
    ...: def func(a, b):
    ...:     return [i for i, x in enumerate(a) if x in b]

In [11]: func(a, b)
   ...: [0, 6, 9]

现在，让我们比较一下这些解决方案的时间性能。

In [12]: %timeit np.where(np.in1d(a, b))[0]
4.26 µs ± 6.9 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

In [13]: %timeit np.nonzero(np.in1d(a, b))[0]
4.39 µs ± 14.3 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

In [14]: %timeit np.searchsorted(a, b)
800 ns ± 6.04 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)

In [15]: %timeit np.searchsorted(a, np.intersect1d(a, b))
8.8 µs ± 73.9 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

In [16]: %timeit [i for i, x in enumerate(a) if x in b]
15.4 µs ± 18.4 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

In [17]: %timeit func(a, b)
336 ns ± 0.579 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)