np.where的两种用法

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)。