假设您有两个二维数组 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 来完成此操作。
最佳答案
方法#1
这是一个基于 View
的。使用 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]))
关于python - 匹配两个二维数组的行并使用 numpy 获取行索引映射,我们在Stack Overflow上找到一个类似的问题: https://stackoverflow.com/questions/55612617/