CODEWITHSUNDEEP
pythonnumpy
Paul Joireman
Paul Joireman

Reputation: 2835

Get array index positions based on conditional

I have a large numpy array and I'd like to get the array indexes based on a given criteria. Numpy provides this but returns a boolean array:

>>> import numpy as np
>>> a = np.arrary([1, 2, 3, 4, 1, 2, 3]
>>> b = a == 3
>>> b
array([False, False, True, False, False, False, True])

but I'd really like to have the actual index positions as integers, is there a simpler way to do that than this:

>>> c = np.arange(len(b))
>>> c = c[b]
>>> c
array([2,6])

In other words, is there a way to do this without creating the c array above?

Upvotes: 1

Views: 283

Answers (2)

Jon Clements
Jon Clements

Reputation: 142206

I'd go for:

import numpy as np
a = np.array([1, 2, 3, 4, 1, 2, 3])
indices, = np.where(a==3)
print indices
# [2 6]

Allows for easier error handling for n-dim arrays (ie, an exception will be thrown if there's too many items to unpack) and it doesn't require flattening.

Upvotes: 3

blazetopher
blazetopher

Reputation: 1050

I believe you're looking for numpy.argwhere:

In[1]: import numpy as np
In[2]: a = np.array([1,2,3,4,1,2,3])
In[3]: b = np.argwhere(a==3).flatten()
In[4]: b
Out[1]: array([2, 6])

Upvotes: 2

Related Questions