Mukundan314
Mukundan314

Reputation: 115

Find index of 2d array in a 3d array in numpy

I want to find a 2d array in a 3d array. for finding a 1d array in a 2d array I can use np.where(np.all(a==b, axis=1))[0][0].

>>> import numpy as np
>>>
>>> a = np.array([[[1, 0, 0],
                   [0, 0, 0],
                   [0, 0, 0]],

                  [[0, 0, 0],
                   [0, 0, 0],
                   [0, 0, 0]]])
>>>
>>> b = np.array([[1, 0, 0],
                  [0, 0, 0],
                  [0, 0, 0]])
>>>
>>> a.find(b)
0

Upvotes: 1

Views: 555

Answers (1)

Paul Panzer
Paul Panzer

Reputation: 53029

The axis keyword accepts tuples, so you can simply do:

np.where(np.all(a==b, axis=(1, 2)))[0][0]

Upvotes: 6

Related Questions