匹配两个二维数组的行并使用 numpy 获取行索引映射答案

【问题标题】：Match rows of two 2D arrays and get a row indices map using numpy匹配两个二维数组的行并使用 numpy 获取行索引映射
【发布时间】：2019-08-31 21:59:15
【问题描述】：

假设您有两个二维数组 A 和 B，并且您想检查 B 中包含 A 的行的位置。你如何使用 numpy 最有效地做到这一点？

例如

a = np.array([[1,2,3],
              [4,5,6],
              [9,10,11]])

b = np.array([[4,5,6],
              [4,3,2],
              [1,2,3],
              [4,8,9]])
map = [[0,2], [1,0]]  # row 0 of a is at row index 2 of array B

我知道如何使用 in1d (test for membership in a 2d numpy array) 检查 A 的行是否在 B 中，但这不会产生索引映射。

此映射的目的是（最终）基于某些列将两个数组合并在一起。
当然可以逐行执行此操作，但这会变得非常低效，因为我的数组具有形状（50 Mio.，20）。

另一种方法是使用pandas merge function，但我想只使用numpy。

【问题讨论】：

数组中的数字是ints吗？如果是这样，它们是积极的吗？如果有，它们是否有可能被限制的范围？
不幸的是，它们是 float64 没有任何限制。
数组中是否存在重复项（相同的行）？
不，在我的情况下不是。

标签： python arrays pandas numpy merge

【解决方案1】：

方法#1

这是基于views 的一个。利用np.argwhere (docs) 返回满足条件的元素的索引，在本例中为成员资格。 -

def view1D(a, b): # a, b are arrays
    a = np.ascontiguousarray(a)
    b = np.ascontiguousarray(b)
    void_dt = np.dtype((np.void, a.dtype.itemsize * a.shape[1]))
    return a.view(void_dt).ravel(),  b.view(void_dt).ravel()

def argwhere_nd(a,b):
    A,B = view1D(a,b)
    return np.argwhere(A[:,None] == B)

方法 #2

这是另一个 O(n)，因此性能要好得多，尤其是在大型阵列上 -

def argwhere_nd_searchsorted(a,b):
    A,B = view1D(a,b)
    sidxB = B.argsort()
    mask = np.isin(A,B)
    cm = A[mask]
    idx0 = np.flatnonzero(mask)
    idx1 = sidxB[np.searchsorted(B,cm, sorter=sidxB)]
    return idx0, idx1 # idx0 : indices in A, idx1 : indices in B

方法#3

另一个O(n) 一个使用argsort() -

def argwhere_nd_argsort(a,b):
    A,B = view1D(a,b)
    c = np.r_[A,B]
    idx = np.argsort(c,kind='mergesort')
    cs = c[idx]
    m0 = cs[:-1] == cs[1:]
    return idx[:-1][m0],idx[1:][m0]-len(A)

使用与之前相同的输入运行示例 -

In [650]: argwhere_nd_searchsorted(a,b)
Out[650]: (array([0, 1]), array([2, 0]))

In [651]: argwhere_nd_argsort(a,b)
Out[651]: (array([0, 1]), array([2, 0]))

【讨论】：

【解决方案2】：

您可以利用自动广播：

np.argwhere(np.all(a.reshape(3,1,-1) == b,2))

导致

array([[0, 2],
       [1, 0]])

注意浮点数，您可能希望将 == 替换为 np.islclose()

【讨论】：