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

Question

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

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

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

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

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