[FIXED] Rowwise numpy.isin for 2D arrays

Issue

I have two arrays:

A = np.array([[3, 1], [4, 1], [1, 4]])
B = np.array([[0, 1, 5], [2, 4, 5], [2, 3, 5]])

Is it possible to use numpy.isin rowwise for 2D arrays? I want to check if A[i,j] is in B[i] and return this result into C[i,j]. At the end I would get the following C:

np.array([[False, True], [True, False], [False, False]])

It would be great, if this is also doable with the == operator, then I could use it also with PyTorch.

Edit: I also considered check for identical rows in different numpy arrays. The question is somehow similar, but I am not able to apply its solutions to this slightly different problem.

Solution

Not sure that my code solves your problem perfectly. please run it on more test cases to confirm. but i would do smth like i've done in the following code taking advantage of numpys vector outer operations ability (similar to vector outer product). If it works as intended it should work with pytorch as well.

import numpy as np
A = np.array([[3, 1], [4, 1], [1, 4]])
B = np.array([[0, 1, 5], [2, 4, 5], [2, 3, 5]])

AA = A.reshape(3, 2, 1)
BB = B.reshape(3, 1, 3)
(AA == BB).sum(-1).astype(bool)

output:

array([[False,  True],
       [ True, False],
       [False, False]])

Answered By - yann ziselman