5.8. Comparisons

We can use comparison operators such as > or == to compare two arrays element-wise. In this case, the result is an array of bool:

import numpy as np

a = np.array(
    [
        [1.0, 2.0, 3.0],
        [4.0, 5.0, 6.0],
        [7.0, 8.0, 9.0],
    ]
)

b = np.array(
    [
        [5.0, 5.0, 5.0],
        [5.0, 5.0, 5.0],
        [5.0, 5.0, 5.0],
    ]
)

a > b

array([[False, False, False],
       [False, False,  True],
       [ True,  True,  True]])

We can also compare arrays with a single number instead of a whole array:

a > 5.0

array([[False, False, False],
       [False, False,  True],
       [ True,  True,  True]])

a == 5.0

array([[False, False, False],
       [False,  True, False],
       [False, False, False]])

a >= 5.0

array([[False, False, False],
       [False,  True,  True],
       [ True,  True,  True]])

To know the indexes of the array where the comparison returned true, we can use np.nonzero:

a = np.array([1.0, 2.2, 10.4, 4.6, 5.8, 6.0, 1.2])

print("Result of the comparison:", a > 5)
print("Indexes where the result was True:", np.nonzero(a > 5))

Result of the comparison: [False False  True False  True  True False]
Indexes where the result was True: (array([2, 4, 5]),)