在 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 上生成相应的布尔向量的明显方法似乎非常慢。

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 方法，将输入转换为 1D 等效项，然后排序并使用 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

示例运行 -

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

Answer 2

numpy：

首先使用np.ravel_multi_index 收集a2 的行和列下标的线性索引等效项。加 1 以说明 numpy 的基于 0 的索引。通过np.unique() 获取存在的唯一行的相应计数。接下来，通过将a1 扩展到向右轴的新维度（也称为广播），在 a2 和 a1 的唯一行之间找到匹配行，并提取非两个数组都为零行。

初始化一个零数组并根据获得的索引通过切片来填充它的值。

def comp2D_vect(a1, a2):
    midx = np.ravel_multi_index(a2.T, a2.max(0)+1)
    a, idx, cnt = np.unique(midx, return_counts=True, return_index=True)
    m1, m2 = (a1[:, None] == a2[idx]).all(-1).nonzero()
    out = np.zeros(a1.shape[0], dtype=int)
    out[m1] = cnt[m2]
    return out

基准测试：

收件人：a2 = a2.repeat(100000, axis=0)

%%timeit
df = pd.DataFrame(a2, columns=['a', 'b', 'c'])
df_count = df.groupby(df.columns.tolist()).size()
df_count.reindex(a1.T.tolist(), fill_value=0).values
10 loops, best of 3: 67.2 ms per loop    # @ Ted Petrou's solution

%timeit comp2D_vect(a1, a2)
10 loops, best of 3: 34 ms per loop      # Posted solution

%timeit searchsorted_based(a1,a2)
10 loops, best of 3: 27.6 ms per loop    # @ Divakar's solution (winner)

Answer 3

Pandas 将是一个很好的工具。您可以将 arr2 放入数据框中，并使用groupby 方法计算每行的出现次数，然后使用 arr1 重新索引结果。

arr1=np.array([[3, 1, 4],
       [4, 3, 5],
       [6, 5, 4],
       [6, 5, 4],
       [3, 1, 4]])

arr2 = np.array([[3, 1, 4],
       [8, 5, 4],
       [3, 1, 4],
       [6, 5, 4],
       [3, 1, 4]])

df = pd.DataFrame(arr2, columns=['a', 'b', 'c'])
df_count = df.groupby(df.columns.tolist()).size()
df_count.reindex(arr1.T.tolist(), fill_value=0)

输出

a  b  c
3  1  4    3
4  3  5    0
6  5  4    1
      4    1
3  1  4    3
dtype: int64

时间安排
首先创建更多数据

arr2_2 = arr2.repeat(100000, axis=0)

现在是时候了：

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

1 个循环，3 个循环中的最佳值：每个循环 704 毫秒

%%timeit
df = pd.DataFrame(arr2_2, columns=['a', 'b', 'c'])
df_count = df.groupby(df.columns.tolist()).size()
df_count.reindex(arr1.T.tolist(), fill_value=0)

10 个循环，3 个循环中的最佳值：每个循环 53.8 毫秒