Question

I need to find the indicies of both the zero and nonzero elements of an array.

Put another way, I want to find the complementary indices from numpy.nonzero().

The way that I know to do this is as follows:

indices_zero = numpy.nonzero(array == 0)
indices_nonzero = numpy.nonzero(array != 0)

This however means searching the array twice, which for large arrays is not efficient. Is there an efficient way to do this using numpy?

Answer 1

Assuming you already have the range for use numpy.arange(len(array)), just get and store the logical indices:

bindices_zero = (array == 0)

then when you actually need the integer indices you can do

indices_zero = numpy.arange(len(array))[bindices_zero]

or

indices_nonzero = numpy.arange(len(array))[~bindices_zero]