How to the function 'where' generate this array?

Question

>>> x = np.arange(9.).reshape(3, 3)
>>> np.where( x > 5 )
(array([2, 2, 2]), array([0, 1, 2]))

What does the x>5 exactly mean? The resulting array seems mysterious.

Answer 1

It's a tuple with row and column indices. x > 5 returns a boolean array of the same shape as x with elements set to True where the condition is fulfilled and False otherwise. According to the documentation np.where will fallback on condition.nonzero when given no other arguments. For your given example all elements greater than 5 happen to be in row 2 and all columns fulfill the condition, hence the [2, 2, 2] (rows), [0, 1, 2] (columns). Note that you can use this result to index the original array:

>>> x[np.where(x > 5)]
[6 7 8]

Answer 2

The usual syntax is np.where(condition, res_if_true, res_if_false). With only the first argument, this is a special case described in the docs:

When only condition is provided, this function is a shorthand for np.asarray(condition).nonzero().

So first calculate x > 5:

arr = x > 5
print(arr)
# array([[False, False, False],
#        [False, False, False],
#        [ True,  True,  True]])

Since it's already an array, calculate arr.nonzero():

print(arr.nonzero())
# (array([2, 2, 2], dtype=int64), array([0, 1, 2], dtype=int64))

This return the indices of the elements that are non-zero. The first element of the tuple represents coordinates of axis=0 and the second element coordinates of axis=1, i.e. all values in the 2nd and final row are greater than 5.