[numpy] the use of numpy.where

Official api introduction

Several usages

1. Enter only condition

import numpy as np

a = np.array([0,5,4,1,9,7])
b = np.where(a>5)
print(b)

The result is as follows

Its result is the subscript of the non-zero element in the original array

2. Input array only

1D array

In the official note, it has been noted that if only condition is entered; then its function is equivalent to
np.asarray(condition).nonzero()

like:

b = np.where([0,0,0,1,1,1,0,0,0])
print(b)

The output is:

2D array

If it is 2D, it returns a tuple, the first value represents the subscript of the 0-dimension of the 2D array, and the second value represents the subscript of the 1-dimension of the 2D array

i = np.array([[False, False],
              [False, True],
              [True, True]])

print(np.where(i))

3. Three parameter input, and each parameter is 1D

it is equivalent to
[xv if c else yv for c, xv, yv in zip(condition, x, y)]

As the following example

import numpy as np

a = np.arange(9)
b = np.where(a<5, a, 0)
print(b)

4. Three parameter input, and each parameter is 2D

If it is 2D, the corresponding elements of the identification are obtained from the corresponding True and False arrays

For example,

b = np.where([[True, False], [True, True]],
         [[1, 2], [3, 4]],
         [[9, 8], [7, 6]])
print(b)

The result is as follows,

explain
the 2D array
[True, False], [True, True]
2 rows and 2 columns to identify:
[True, False
True, True]
It identifies the result. True is obtained from the second parameter array, and False is obtained from the second parameter array Obtained from the third parameter array.

---

Summarize

1 parameter outputs the subscript; 3 parameters output the corresponding value