Getting indices of both zero and nonzero elements in array

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]

Answer 2

You can use boolean indexing:

In [82]: a = np.random.randint(-5, 5, 100)

In [83]: a
Out[83]: 
array([-2, -1,  4, -3,  1, -2,  2, -1,  2, -1, -3,  3, -3, -4,  1,  2,  1,
        3,  3,  0,  1, -3, -4,  3, -5, -1,  3,  2,  3,  0, -5,  4,  3, -5,
       -3,  1, -1,  0, -4,  0,  1, -5, -5, -1,  3, -2, -5, -5,  1,  0, -1,
        1,  1, -1, -2, -2,  1,  1, -4, -4,  1, -3, -3, -5,  3,  0, -5, -2,
       -2,  4,  1, -4, -5, -1,  3, -3,  2,  4, -4,  4,  2, -2, -4,  3,  4,
       -2, -4,  2, -4, -1,  0, -3, -1,  2,  3,  1,  1,  2,  1,  4])

In [84]: mask = a != 0

In [85]: a[mask]
Out[85]: 
array([-2, -1,  4, -3,  1, -2,  2, -1,  2, -1, -3,  3, -3, -4,  1,  2,  1,
        3,  3,  1, -3, -4,  3, -5, -1,  3,  2,  3, -5,  4,  3, -5, -3,  1,
       -1, -4,  1, -5, -5, -1,  3, -2, -5, -5,  1, -1,  1,  1, -1, -2, -2,
        1,  1, -4, -4,  1, -3, -3, -5,  3, -5, -2, -2,  4,  1, -4, -5, -1,
        3, -3,  2,  4, -4,  4,  2, -2, -4,  3,  4, -2, -4,  2, -4, -1, -3,
       -1,  2,  3,  1,  1,  2,  1,  4])

In [86]: a[-mask]
Out[86]: array([0, 0, 0, 0, 0, 0, 0])

Answer 3

I'm not sure about a built-in numpy method for accomplishing this, but you could use an old-fashioned for loop, I believe. Something like:

indices_zero = []
indices_nonzero = []

for index in xrange(len(array)):
    if array[index] == 0:
        indicies_zero.append(index)
    else:
        indicies_nonzero.append(index)

Something like this should accomplish what you want, by only looping once.