numpy.where 与 numpy.nonzero

一，numpy.where 接口

注意函数的定义为：numpy.where(condition, elem_x, elem_y)

其中的 condition 可以理解为一个【布尔索引】（对布尔索引不熟悉的可以看我之前写过的文章，或者自己上网查）。对于布尔索引中的元素，如果其值为 True，则将其替换为 elem_x；如果其值为 False，则将其替换为 elem_y。最后将替换后的结果返回。

# 导入 numpy 库
import numpy as np
 
# 创建一个 ndarray 类型数组
arr = np.array(
        [[3, 2, 3, 4, 9],
         [9, 7, 8, 0, 7],
         [1, 3, 3, 4, 5],
         [7, 8, 0, 9, 3],
         [9, 2, 7, 3, 5]
         ])

# 调用接口
out = np.where(arr < 3, 'A', 'B')

#打印结果
print(out)

运行结果：

[['B' 'A' 'B' 'B' 'B']
 ['B' 'B' 'B' 'A' 'B']
 ['A' 'B' 'B' 'B' 'B']
 ['B' 'B' 'A' 'B' 'B']
 ['B' 'A' 'B' 'B' 'B']]

可见，numpy.where 接口会返回一个与 arr 相同规模的 ndarray，其中满足条件的元素会被替换为 'A'，不满足条件的元素会被替换为 'B'。

注意：参数 elem_x 与 elem_y 可以省略。但这两个参数要么都有，要么都没有，否则会报错。

二，numpy.nonzero 接口

当参数 elem_x 与 elem_y 都省略时，numpy.where 接口退化为 numpy.nonzero 接口。根据官方文档的说法，此时应该优先考虑使用 numpy.nonzero 接口。

Note

When only condition is provided, this function is a shorthand for np.asarray(condition).nonzero(). Using nonzero directly should be preferred, as it behaves correctly for subclasses. The rest of this documentation covers only the case where all three arguments are provided.

numpy.nonzero 接口的输入参数是一个 ndarray 数组 arr。返回结果是 arr 中所有不为 0 的元素的下标。因此，而当我们将 condition 作为参数时，会返回布尔索引中所有值为 True 的元素的下标。

# 导入 numpy 库
import numpy as np
 
# 创建一个 ndarray 类型数组
arr = np.array(
        [[3, 2, 3, 4, 9],
         [9, 7, 8, 0, 7],
         [1, 3, 3, 4, 5],
         [7, 8, 0, 9, 3],
         [9, 2, 7, 3, 5]
         ])

# 调用接口
out = np.nonzero(arr < 3)

#打印结果
print(out)

运行结果：

(array([0, 1, 2, 3, 4], dtype=int64), array([1, 3, 0, 2, 1], dtype=int64))

其中，返回值 out 是一个元组。元组的第一个分量表示行坐标 row，第二个分量表示列坐标 col。这两个分量组合成为的索引，指向 arr 中所有满足 arr < 3 条件的元素。

其原理就是：条件表达式 arr < 3 的返回值是一个【布尔索引】，布尔值 True 被解释为非零元素，而布尔值 False 被解释为 0 。