get indicies of non-zero elements of 2D array

Question

From Getting indices of both zero and nonzero elements in array , I can get indicies of non-zero elements in a 1 D array in numpy like this:

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

Is there a way to extend it to a 2D array?

Answer 1

You can use numpy.nonzero

The following code is self-explanatory

import numpy as np

A = np.array([[1, 0, 1],
              [0, 5, 1],
              [3, 0, 0]])
nonzero = np.nonzero(A)
# Returns a tuple of (nonzero_row_index, nonzero_col_index)
# That is (array([0, 0, 1, 1, 2]), array([0, 2, 1, 2, 0]))

nonzero_row = nonzero[0]
nonzero_col = nonzero[1]

for row, col in zip(nonzero_row, nonzero_col):
    print("A[{}, {}] = {}".format(row, col, A[row, col]))
"""
A[0, 0] = 1
A[0, 2] = 1
A[1, 1] = 5
A[1, 2] = 1
A[2, 0] = 3
"""

You can even do this

A[nonzero] = -100
print(A)
"""
[[-100    0 -100]
 [   0 -100 -100]
 [-100    0    0]]
 """

Other variations

np.where(array)

It is equivalent to np.nonzero(array) But, np.nonzero is preferred because its name is clear

np.argwhere(array)

It's equivalent to np.transpose(np.nonzero(array))

print(np.argwhere(A))
"""
[[0 0]
 [0 2]
 [1 1]
 [1 2]
 [2 0]]
 """

Answer 2

A = np.array([[1, 0, 1],
              [0, 5, 1],
              [3, 0, 0]])

np.stack(np.nonzero(A), axis=-1)

array([[0, 0],
       [0, 2],
       [1, 1],
       [1, 2],
       [2, 0]])

np.nonzero returns a tuple of arrays, one for each dimension of a, containing the indices of the non-zero elements in that dimension.

https://docs.scipy.org/doc/numpy/reference/generated/numpy.nonzero.html

np.stack joins this tuple array along a new axis. In our case, the innermost axis also known as the last axis (denoted by -1).

The axis parameter specifies the index of the new axis in the dimensions of the result. For example, if axis=0 it will be the first dimension and if axis=-1 it will be the last dimension.

New in version 1.10.0.

https://docs.scipy.org/doc/numpy/reference/generated/numpy.stack.html