1. np.where 函数可以根据指定的条件返回两个数组中对应位置的元素。
np.where(condition, x, y)
其中,condition 是一个条件表达式,x 和 y 是两个具有相同形状的数组。当 condition 的某个元素为 True 时,返回 x 中对应位置的元素,否则返回 y 中对应位置的元素。
例如,以下代码将数组 a 中大于 0 的元素设置为 1,小于 0 的元素设置为 -1,等于 0 的元素设置为 0:
import numpy as np
a = np.array([-1, 0, 1, -2, 3, 0])
b = np.where(a > 0, 1, np.where(a < 0, -1, 0))
print(b) # 输出 [-1 0 1 -1 1 0]
其中,第一个 np.where 函数用于判断大于 0 的元素,第二个 np.where 函数用于判断小于 0 的元素。先判断第一个np.where(),其中小于0的数进行第二个np.where()的判断。
2. np.where(matrix) 返回一个包含所有非零元素下标的元组
假设矩阵维度为 (row_indices, column_indices)。其中 row_indices 和 column_indices 分别是非零元素的行和列下标组成的一维数组。
举个例子,假设有如下二维数组:
[[0, 1, 0],
[1, 0, 1],
[0, 1, 0]]
调用 np.where() 函数将返回如下元组:
(array([0, 1, 1, 2]), array([1, 0, 2, 1]))
这表示非零元素的下标为 (0, 1), (1, 0), (1, 2), (2, 1)。