Selecting indexes of elements with a common property in Python

Question

I have a numpy array and would like to obtain the indexes of the elements that verify a common property. For example, suppose the array is np.array([1, 0, 0, 0, 1, 1, 0, 0, 1, 0, 1, 0, 0, 0, 1]) , and I want to have the indexes of all elements equal to 1, so the output would be [0, 4, 5, 8, 10, 14] .

I have defined the following procedure

def find_indexes(A):
    res = []
    for i in range(len(A)):
        if A[i] == 1:
            res.append(i)
    return res

Is there a more "pythonesque" way of doing this? More specifically, I am wondering if there is something similar to boolean indexing:

A[A>=1]

that would return the indexes of the elements rather than the elements themselves.

Answer 1

use np.where .

  import numpy as np
  x = np.array(np.array([1, 0, 0, 0, 1, 1, 0, 0, 1, 0, 1, 0, 0, 0, 1])
  indices, = np.where(x == 1)
  print(indices)

Answer 2

Use numpy.where

arr = np.array([1, 0, 0, 0, 1, 1, 0, 0, 1, 0, 1, 0, 0, 0, 1])
print np.where(arr == 1)
(array([ 0,  4,  5,  8, 10, 14]),)

Answer 3

List comprehension for pure python:

ar = [i for i in range(len(a)) if a[i] == 1]