查找两个二维数组中匹配行的索引

时间:2013-11-26 23:47:36

标签: python numpy

假设我有两个2-D数组,如下所示:

array([[3, 3, 1, 0],
       [2, 3, 1, 3],
       [0, 2, 3, 1],
       [1, 0, 2, 3],
       [3, 1, 0, 2]], dtype=int8)

array([[0, 3, 3, 1],
       [0, 2, 3, 1],
       [1, 0, 2, 3],
       [3, 1, 0, 2],
       [3, 3, 1, 0]], dtype=int8)

每个数组中的某些行都有一个对应的行,它在另一个数组中按值(但不一定是索引)匹配,有些则没有。

我想找到一种有效的方法来返回两个与匹配行对应的数组中的索引对。如果它们是元组我会期望返回

(0,4)
(2,1)
(3,2)
(4,3)

3 个答案:

答案 0 :(得分:7)

这是一个全numpy解决方案 - 不一定比迭代Python更好。它仍然需要考虑所有组合。

In [53]: np.array(np.all((x[:,None,:]==y[None,:,:]),axis=-1).nonzero()).T.tolist()
Out[53]: [[0, 4], [2, 1], [3, 2], [4, 3]]

中间数组是(5,5,4)np.all将其缩小为:

array([[False, False, False, False,  True],
       [False, False, False, False, False],
       [False,  True, False, False, False],
       [False, False,  True, False, False],
       [False, False, False,  True, False]], dtype=bool)

其余的只是提取指数True

在原油测试中,这次是47.8 us;另一个答案是{38}我们的L1字典;和第三个在496我们的双循环。

答案 1 :(得分:5)

我想不出一个特定的方式来做这件事,但这就是我对常规名单所做的事情:

>>> L1= [[3, 3, 1, 0],
...        [2, 3, 1, 3],
...        [0, 2, 3, 1],
...        [1, 0, 2, 3],
...        [3, 1, 0, 2]]
>>> L2 = [[0, 3, 3, 1],
...        [0, 2, 3, 1],
...        [1, 0, 2, 3],
...        [3, 1, 0, 2],
...        [3, 3, 1, 0]]
>>> L1 = {tuple(row):i for i,row in enumerate(L1)}
>>> answer = []
>>> for i,row in enumerate(L2):
...   if tuple(row) in L1:
...     answer.append((L1[tuple(row)], i))
... 
>>> answer
[(2, 1), (3, 2), (4, 3), (0, 4)]

答案 2 :(得分:5)

您可以使用void数据类型技巧在两个数组的行上使用1D函数。 a_viewb_view是1D向量,每个条目代表一个完整的行。然后我选择对数组进行排序并使用np.searchsorted来查找该数组中另一个数组的项。如果我们排序的数组的长度为m而另一个数组的长度为n,则排序需要时间m * log(m),而np.searchsorted的二进制搜索需要时间n * log(m) ,总共(n + m) * log(m)。因此,您希望对两个数组中最短的数组进行排序:

def find_rows(a, b):
    dt = np.dtype((np.void, a.dtype.itemsize * a.shape[1]))

    a_view = np.ascontiguousarray(a).view(dt).ravel()
    b_view = np.ascontiguousarray(b).view(dt).ravel()

    sort_b = np.argsort(b_view)
    where_in_b = np.searchsorted(b_view, a_view,
                                 sorter=sort_b)
    where_in_b = np.take(sort_b, where_in_b)
    which_in_a = np.take(b_view, where_in_b) == a_view
    where_in_b = where_in_b[which_in_a]
    which_in_a = np.nonzero(which_in_a)[0]
    return np.column_stack((which_in_a, where_in_b))

使用ab两个示例数组:

In [14]: find_rows(a, b)
Out[14]: 
array([[0, 4],
       [2, 1],
       [3, 2],
       [4, 3]], dtype=int64)

In [15]: %timeit find_rows(a, b)
10000 loops, best of 3: 29.7 us per loop

在我的系统上,字典方法在测试数据的时间约为22 us时更快,但是对于1000x4的数组,这种numpy方法比纯Python方法(483 us vs 2.54 ms)快约6倍。