Approach #1
Here's one based on views
. Makes use of np.argwhere
(docs) to return the indices of an element that meet a condition, in this case, membership. -
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)
Approach #2
Here's another that would be O(n)
and hence much better on performance, especially on large arrays -
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
Approach #3
Another O(n)
one using 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)
Sample runs with same inputs as earlier -
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]))