如果大数组中包含小数组的值，找到它们的索引。

Question

有没有一个快速的numpy函数，可以返回一个较大数组中与较小数组匹配的值的索引列表？较小的数组有约30M个值，而较大的数组有8亿个值，因此我想避免使用numpy.where循环调用。

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

不要这样做：

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

I would want this:

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

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

- isosceleswheel

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)

对于大的a和b，使用searchsorted可能比np.in1d(a, b)更快：

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

- unutbu

---

使用searchsorted的左右智能操作！不过，猜测你需要使用掩码来索引np.searchsorted(b, a, side='left')以获取实际的索引。 - Divakar

你仍需要使用 np.where 来获取相对于 a 的索引。 - unutbu

我的意思是这样的：idx1 = np.searchsorted(a,b,'left'); idx2 = np.searchsorted(a,b,'right'); out = idx1[idx1 != idx2]。也许可以工作？ - Divakar

@Divakar: 看起来比我建议的解决方案快了大约5倍（对于较大的测试用例）。你想写出来吗? - unutbu

感谢您的动力！我必须说，很大程度上受到了您解决方案的启发。刚刚在这里发布了解决方案。 - Divakar

谢谢你们的努力，这对我来说已经足够了！ - isosceleswheel

Answer 2

使用np.searchsorted时，默认的搜索方向是left。我们也可以从right方向进行搜索，并且在这两组索引中相同的内容应该避免从left选项输出的索引中获取以获得所需的输出。这里的动机与@unutbu's solution中讨论的相同。

因此，实现代码如下 -

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

- Divakar

---

ن½؟用searchsorted(larger, smaller)ن»£و›؟searchsorted(smaller, larger)ه¹¶ç´¢ه¼•idx1而ن¸چوک¯ن½؟用np.where，è؟™وک¯ن¸€ç§چو›´هٹ èپھوکژçڑ„و–¹و³•م€‚ - unutbu

您需要添加额外的处理，因为 searchsorted 可能返回超出界限的索引，但是一种可能更快地丢弃非匹配项的方法，而不是使用“right”重新搜索，是检查它们是否匹配！ idx = np.searchsorted(haystack, needle); mask = idx < haystack.size; mask[mask] = haystack[idx[mask]] == needle[mask]; idx = idx[mask] - Jaime