Use of loops (example, for) in the conditional arguments of numpy.where-CodePudding

I have a matrix "x" and an array "index". I just want to find the position (row number) of the "index" array in the matrix "x". Example:

x = np.array([[0, 0],
       [1, 0],
       [2, 0],
       [0, 1],
       [1, 1],
       [2, 1],
       [0, 2],
       [1, 2],
       [2, 2],
       [0, 3],
       [1, 3],
       [2, 3],])

index = [2,1]

Here if I use the code: np.where((x[:,0]==index[0]) & (x[:,1]==index[1]))[0] it is working.

But if I have the matrix "x" with N number of columns (instead of 2), I have to use loop inside the np.where arguments. I tried this: np.where((for b in range(2):(x[:,b]==index[b]) & (x[:,b]==index[b])))[0]

Then it shows "invalid syntax" error. Can you please help me regarding this? Thanks in advance.

CodePudding user response：

You can use broadcasting:

np.where((x == np.array(index).reshape(1, -1)).all(1))

or even:

np.where((x == np.array(index)[None, :]).all(1))

both results in:

(array([5], dtype=int64),)

CodePudding user response：

if index be an array (as it is written in the question, not in the code):

x = np.array([[0, 0, 7],
              [1, 0, 8],
              [2, 3, 11],
              [0, 1, 1],
              [1, 1, 9],
              [2, 1, 4],      # <--
              [0, 2, 7],
              [1, 2, 3],
              [2, 1, 4],      # <--
              [2, 2, 1],
              [0, 3, 2],
              [1, 3, 11],     # <--
              [1, 3, 8],
              [2, 3, 8],])

index = np.array([[2, 1, 4], [1, 3, 11]])

it can be done by:

np.where((x == index[:, None]).all(-1))[1]
# [ 5  8 11]