CODEWITHSUNDEEP
pythonnumpy
Neil G
Neil G

Reputation: 33202

Python numpy index set from Boolean array

How do I transform a Boolean array into an iterable of indexes?

E.g.,

import numpy as np
import itertools as it
x = np.array([1,0,1,1,0,0])
y = x > 0
retval = [i for i, y_i in enumerate(y) if y_i]

Is there a nicer way?

Upvotes: 2

Views: 1278

Answers (1)

mathematical.coffee
mathematical.coffee

Reputation: 56905

Try np.where or np.nonzero.

x = np.array([1, 0, 1, 1, 0, 0])
np.where(x)[0] # returns a tuple hence the [0], see help(np.where)
# array([0, 2, 3])
x.nonzero()[0] # in this case, the same as above.

See help(np.where) and help(np.nonzero).

Possibly worth noting that in the np.where page it mentions that for 1D x it's basically equivalent to your longform in the question.

Upvotes: 3

Related Questions