如果大数组包含较小数组中的值，则查找大数组的索引

Question

是否有一个快速的 numpy 函数用于返回较大数组中的索引列表，它与较小数组中的值匹配？较小的数组约为 30M 值，较大的是 800M，所以我想避免 numpy.where 调用的 for 循环。

searchsorted 的问题在于，即使它们不是完全匹配，它也会返回结果，它只是给出最接近的索引，但我只想要完全匹配的索引

而不是这个：

>>> a = array([1,2,3,4,5])
>>> b = array([2,4,7])
>>> searchsorted(a,b)
array([1, 3, 5])

我想要这个：

>>> a = array([1,2,3,4,5])
>>> b = array([2,4,7])
>>> SOMEFUNCTION(a,b)
array([1, 3])

编辑：较小和较大数组中的值集始终是唯一且已排序的。

Answer 1

您可以使用np.in1d 来查找a 中位于b 中的那些元素。 要查找索引，请调用np.where：

In [34]: a = array([1,2,3,4,5])

In [35]: b = array([2,4,7])

In [36]: np.in1d(a, b)
Out[38]: array([False,  True, False,  True, False], dtype=bool)

In [39]: np.where(np.in1d(a, b))
Out[39]: (array([1, 3]),)

因为a 和b 已经排序，你可以使用

In [57]: np.searchsorted(b, a, side='right') != np.searchsorted(b, a, side='left')
Out[57]: array([False,  True, False,  True, False], dtype=bool)

而不是np.in1d(a, b)。对于较大的a 和b，使用searchsorted 可能会更快：

import numpy as np
a = np.random.choice(10**7, size=10**6, replace=False)
a.sort()
b = np.random.choice(10**7, size=10**5, replace=False)
b.sort()

In [53]: %timeit np.in1d(a, b)
10 loops, best of 3: 176 ms per loop

In [54]: %timeit np.searchsorted(b, a, side='right') != np.searchsorted(b, a, side='left')
10 loops, best of 3: 106 ms per loop

---

Jaime 和Divakar 建议对上述方法进行一些重大改进。下面是一些测试方法是否都返回相同结果的代码，然后是一些基准测试：

import numpy as np

a = np.random.choice(10**7, size=10**6, replace=False)
a.sort()
b = np.random.choice(10**7, size=10**5, replace=False)
b.sort()

def using_searchsorted(a, b):
    return (np.where(np.searchsorted(b, a, side='right') 
                     != np.searchsorted(b, a, side='left')))[0]

def using_in1d(a, b):
    return np.where(np.in1d(a, b))[0]

def using_searchsorted_divakar(a, b):
    idx1 = np.searchsorted(a,b,'left')
    idx2 = np.searchsorted(a,b,'right')
    out = idx1[idx1 != idx2]
    return out

def using_jaime_mask(haystack, needle):
    idx = np.searchsorted(haystack, needle)
    mask = idx < haystack.size
    mask[mask] = haystack[idx[mask]] == needle[mask]
    idx = idx[mask]
    return idx

expected = using_searchsorted(a, b)
for func in (using_in1d, using_searchsorted_divakar, using_jaime_mask):
    result = func(a, b)
    assert np.allclose(expected, result)

---

In [29]: %timeit using_jaime_mask(a, b)
100 loops, best of 3: 13 ms per loop

In [28]: %timeit using_searchsorted_divakar(a, b)
10 loops, best of 3: 21.7 ms per loop

In [26]: %timeit using_searchsorted(a, b)
10 loops, best of 3: 109 ms per loop

In [27]: %timeit using_in1d(a, b)
10 loops, best of 3: 173 ms per loop

Answer 2

np.searchsorted 的默认seacrh 方向是left。我们也可以从right 方向搜索它，并且在这两个索引集中相同的索引将是从left 选项输出的索引中要避免的索引以获得所需的输出。这里的动机与@unutbu's solution 中讨论的相同。

因此，实现看起来像这样 -

idx1 = np.searchsorted(a,b,'left')
idx2 = np.searchsorted(a,b,'right')
out = idx1[idx1 != idx2]