python - 在 NumPy 中按行比较两个矩阵

Question

假设我有两个 NumPy 矩阵(或 Pandas DataFrames，但我猜这在 NumPy 中会更快)。

>>> arr1
array([[3, 1, 4],
       [4, 3, 5],
       [6, 5, 4],
       [6, 5, 4],
       [3, 1, 4]])
>>> arr2
array([[3, 1, 4],
       [8, 5, 4],
       [3, 1, 4],
       [6, 5, 4],
       [3, 1, 4]])

对于 arr1 中的每个行向量，我想计算该行向量在 arr2 中的出现次数，并生成这些计数的向量。所以对于这个例子，结果将是

[3, 0, 1, 1, 3]

执行此操作的有效方法是什么？

第一种方法:仅使用遍历 arr1 的行向量并在 arr2 上生成相应的 bool 向量的明显方法似乎非常慢。

np.apply_along_axis(lambda x: (x == arr2).all(1).sum(), axis=1, arr=arr1)

这似乎是一个糟糕的算法，因为我必须多次检查相同的行。

第二种方法:我可以将行计数存储在 collections.Counter 中，然后使用 apply_along_axis 访问它。

cnter = Counter(tuple(row) for row in arr2)
np.apply_along_axis(lambda x: cnter[tuple(x)], axis=1, arr=arr1)

这似乎更快一些，但我觉得必须有比这更直接的方法。

Answer 1

这是一种 NumPy 方法，将输入转换为一维等价物，然后排序并使用 np.searchsorted 和 np.bincount 进行计数 -

def searchsorted_based(a,b):      
    dims = np.maximum(a.max(0), b.max(0))+1

    a1D = np.ravel_multi_index(a.T,dims)
    b1D = np.ravel_multi_index(b.T,dims)

    unq_a1D, IDs = np.unique(a1D, return_inverse=1)
    fidx = np.searchsorted(unq_a1D, b1D)
    fidx[fidx==unq_a1D.size] = 0
    mask = unq_a1D[fidx] == b1D 

    count = np.bincount(fidx[mask])
    out = count[IDs]
    return out

sample 运行-

In [308]: a
Out[308]: 
array([[3, 1, 4],
       [4, 3, 5],
       [6, 5, 4],
       [6, 5, 4],
       [3, 1, 4]])

In [309]: b
Out[309]: 
array([[3, 1, 4],
       [8, 5, 4],
       [3, 1, 4],
       [6, 5, 4],
       [3, 1, 4],
       [2, 1, 5]])

In [310]: searchsorted_based(a,b)
Out[310]: array([3, 0, 1, 1, 3])

运行时测试-

In [377]: A = a[np.random.randint(0,a.shape[0],(1000))]

In [378]: B = b[np.random.randint(0,b.shape[0],(1000))]

In [379]: np.allclose(comp2D_vect(A,B), searchsorted_based(A,B))
Out[379]: True

# @Nickil Maveli's soln
In [380]: %timeit comp2D_vect(A,B)
10000 loops, best of 3: 184 µs per loop

In [381]: %timeit searchsorted_based(A,B)
10000 loops, best of 3: 92.6 µs per loop