CODEWITHSUNDEEP
numpy
Henrik
Henrik

Reputation: 89

How to find matrix common members of matrices in Numpy

I have a 2D matrix A and a vector B. I want to find all row indices of elements in A that are also contained in B.

A = np.array([[1,9,5], [8,4,9], [4,9,3], [6,7,5]], dtype=int)
B = np.array([2, 4, 8, 10, 12, 18], dtype=int)

My current solution is only to compare A to one element of B at a time but that is horribly slow:

res = np.array([], dtype=int)
for i in range(B.shape[0]):
    cres, _ = (B[i] == A).nonzero()
    degElem = np.append(res, cres)
res = np.unique(res)

The following Matlab statement would solve my issue:

find(any(reshape(any(reshape(A, prod(size(A)), 1) == B, 2),size(A, 1),size(A, 2)), 2))

However comparing a row and a colum vector in Numpy does not create a Boolean intersection matrix as it does in Matlab. Is there a proper way to do this in Numpy?

Upvotes: 1

Views: 52

Answers (1)

Divakar
Divakar

Reputation: 221504

We can use np.isin masking.

To get all the row numbers, it would be -

np.where(np.isin(A,B).T)[1]

If you need them split based on each element's occurence -

[np.flatnonzero(i) for i in np.isin(A,B).T if i.any()]

Posted MATLAB code seems to be doing broadcasting. So, an equivalent one would be -

np.where(B[:,None,None]==A)[1]

Upvotes: 1

Related Questions