在两个数组中找到共同值的索引

Question

在两个数组中找到共同值的索引

5

我正在使用Python 2.7。我有两个数组A和B。要查找在B中存在的元素在A中的索引，我可以这样做：

A_inds = np.in1d(A,B)

我也想获取在A中存在的元素在B中的索引，即使用上述代码找到的相同重叠元素的B中的索引。

目前，我正在以以下方式再次运行同一行：

B_inds = np.in1d(B,A)

但是这个额外的计算似乎是不必要的。有没有更加计算有效的方法来获取A_inds和B_inds？我可以使用列表或数组方法。

- berkelem

输入数组的大小是多少？它们是一维的吗？ - Divakar

大型的。数量级在10^6或10^7左右。 - berkelem

1

这些数组元素是唯一的吗？它们是否已排序？ - Divakar

很遗憾，不是的。这个数组中有许多重复元素，大约占5-10%。而且，它们都是一维的。 - berkelem

如果有更好的独特元素解决方案，我可能会调整一下代码使它们变得独特。目前，重复出现是由于代码中先前的四舍五入而引起的。 - berkelem

1

元素并不是严格排序的。实际上，它们是元组。也许我应该早点提到这一点。 - berkelem

2个回答

1

一个简单的多进程实现将使您获得更快的速度：

import time
import numpy as np

from multiprocessing import Process, Queue

a = np.random.randint(0, 20, 1000000)
b = np.random.randint(0, 20, 1000000)

def original(a, b, q):
    q.put( np.in1d(a, b) )

if __name__ == '__main__':
    t0 = time.time()
    q = Queue()
    q2 = Queue()
    p = Process(target=original, args=(a, b, q,))
    p2 = Process(target=original, args=(b, a, q2))
    p.start()
    p2.start()
    res = q.get()
    res2 = q2.get()

    print time.time() - t0

>>> 0.21398806572

Divakar的unq_searchsorted(A,B)方法在我的机器上花费了0.271834135056秒。

- kezzos

谢谢您提供这个 - 它肯定会很有用。不过现在，我正在寻找单核心上最快的方法，因为稍后我将把整个代码分发到几个核心上。 - berkelem

网页内容由stack overflow 提供, 点击上面的

可以查看英文原文，
原文链接

- Divakar · Accepted Answer

np.unique和np.searchsorted可以一起使用来解决这个问题 -

def unq_searchsorted(A,B):

    # Get unique elements of A and B and the indices based on the uniqueness
    unqA,idx1 = np.unique(A,return_inverse=True)
    unqB,idx2 = np.unique(B,return_inverse=True)

    # Create mask equivalent to np.in1d(A,B) and np.in1d(B,A) for unique elements
    mask1 = (np.searchsorted(unqB,unqA,'right') - np.searchsorted(unqB,unqA,'left'))==1
    mask2 = (np.searchsorted(unqA,unqB,'right') - np.searchsorted(unqA,unqB,'left'))==1

    # Map back to all non-unique indices to get equivalent of np.in1d(A,B), 
    # np.in1d(B,A) results for non-unique elements
    return mask1[idx1],mask2[idx2]

运行时测试和验证结果 -

In [233]: def org_app(A,B):
     ...:     return np.in1d(A,B), np.in1d(B,A)
     ...: 

In [234]: A = np.random.randint(0,10000,(10000))
     ...: B = np.random.randint(0,10000,(10000))
     ...: 

In [235]: np.allclose(org_app(A,B)[0],unq_searchsorted(A,B)[0])
Out[235]: True

In [236]: np.allclose(org_app(A,B)[1],unq_searchsorted(A,B)[1])
Out[236]: True

In [237]: %timeit org_app(A,B)
100 loops, best of 3: 7.69 ms per loop

In [238]: %timeit unq_searchsorted(A,B)
100 loops, best of 3: 5.56 ms per loop

如果两个输入数组已经是排序且唯一的，性能提升将非常大。因此，解决方案函数将简化为 -

def unq_searchsorted_v1(A,B):
    out1 = (np.searchsorted(B,A,'right') - np.searchsorted(B,A,'left'))==1
    out2 = (np.searchsorted(A,B,'right') - np.searchsorted(A,B,'left'))==1  
    return out1,out2

后续的运行时测试 -

In [275]: A = np.random.randint(0,100000,(20000))
     ...: B = np.random.randint(0,100000,(20000))
     ...: A = np.unique(A)
     ...: B = np.unique(B)
     ...: 

In [276]: np.allclose(org_app(A,B)[0],unq_searchsorted_v1(A,B)[0])
Out[276]: True

In [277]: np.allclose(org_app(A,B)[1],unq_searchsorted_v1(A,B)[1])
Out[277]: True

In [278]: %timeit org_app(A,B)
100 loops, best of 3: 8.83 ms per loop

In [279]: %timeit unq_searchsorted_v1(A,B)
100 loops, best of 3: 4.94 ms per loop