Introduction to NumPy and Matplotlib
Chapter 6: NumPy functions
Searching functions
We have previously seen how boolean expressions can be used to filter a list.
>>> x = np.array([[0, 2, 2], [0, 3, 0]])
>>> print(x != 0)
[[False True True]
[False True False]]
>>> print(x[x != 0])
[2 2 3]
What if we want to know where in the array is x
not equal to zero? There is a NumPy
function for that!
np.nonzero(x)
or x.nonzero()
returns the indices of all elements in x
that are not zero (or not False
).
>>> x = np.array([[0, 2, 2], [0, 3, 0]])
>>> print(np.nonzero(x))
(array([0, 0, 1], dtype=int64), array([1, 2, 1], dtype=int64))
>>> print(x[np.nonzero(x)]
[2 2 3]
Again, you should think of the indices returned by np.nonzero(x)
sort of like the output of Python’s zip()
. Here is a diagram which explains this better: