CODEWITHSUNDEEP
pythonarraysnumpy
Carven
Carven

Reputation: 15660

How to find the index of array by a value which is an array?

I'm trying to find the index of a numpy array by value. The value however is also an array. In other words, it's a multi dimensional array.

For example:

a = [
[[1, 0], [0, 2], [3, 3]], 
[[1, 0], [1, 3], [1, 0]], 
[[4, 0], [2, 3], [3, 0]]
]

I want to find the index of [[1, 0], [1, 3], [1, 0]], which is 1. Basically, I want to find the element in the array which matches the array pattern that I have.

How can I do this using numpy?

Upvotes: 0

Views: 73

Answers (2)

cs95
cs95

Reputation: 403010

Use np.flatnonzero in conjunction with broadcasted comparison:

a

array([[[1, 0],
        [0, 2],
        [3, 3]],

       [[1, 0],
        [1, 3],
        [1, 0]],

       [[4, 0],
        [2, 3],
        [3, 0]]])

np.flatnonzero((a == [[1, 0], [1, 3], [1, 0]]).all(1).all(1))
array([1])

Borrowing from the other answer, you could pass multiple axes to all:

np.flatnonzero((a == [[1, 0], [1, 3], [1, 0]]).all((1, 2)))
array([1])

Upvotes: 2

Paul Panzer
Paul Panzer

Reputation: 53089

You can use np.all and np.where:

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

Upvotes: 1

Related Questions