numpy where compare arrays as a whole

Question

I have an array x=np.array([[0,1,2,],[0,0,0],[3,4,0],[1,2,3]]), and I want to get the index where x=[0,0,0], i.e. 1. I tried np.where(x==[0,0,0]) resulting in (array([0, 1, 1, 1, 2]), array([0, 0, 1, 2, 2])). How can I get the desired answer?

Answer 1

As @transcranial solution, you can use np.all() to do the job. But np.all() is slow, so if you apply it to a large array, speed will be your concern.

To test for a specific value or a specific range, I would do like this.

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

condition = (x[:,0]==0) & (x[:,1]==0) & (x[:,2]==0)
np.where(condition)
# (array([1, 4]),)

It's a bit ugly but it almost twice as fast as np.all() solution.

In[23]: %timeit np.where(np.all(x == [0,0,0], axis=1) == True)
100000 loops, best of 3: 6.5 µs per loop
In[22]: %timeit np.where((x[:,0]==0)&(x[:,1]==0)&(x[:,2]==0))
100000 loops, best of 3: 3.57 µs per loop

And you can test not only for equality but also a range.

condition = (x[:,0]<3) & (x[:,1]>=1) & (x[:,2]>=0)
np.where(condition)
# (array([0, 3]),)

Answer 2

Bit expensive likely but x.tolist().index([0,0,0]) should work

To adapt to get all matching indices instead of one, use

xL = x.tolist()
indices = [i for i, x in enumerate(xL) if x == [0,0,0]]

(inspired by this answer)

Answer 3

import numpy as np
x = np.array([[0,1,2],[0,0,0],[3,4,0],[1,2,3]])
np.where(np.all(x == [0,0,0], axis=1) == True)[0].tolist()
# outputs [1]

np.all tests for equality for all elements along axis=1, which outputs array([False, True, False, False]) in this case. Then use np.where to get the index where this is true.