NumPy np.where() 按条件替换值

说明

NumPy 教程持续更新中，提供建议、纠错、催更等加作者微信: gairuo123（备注：pandas教程）和关注公众号「盖若」ID: gairuo。跟作者学习，请进入 Python学习课程。欢迎关注作者出版的书籍：《深入浅出Pandas》和《Python之光》。

np.where() 是 NumPy 中的一个非常有用的函数，主要用于根据条件选择元素或者进行条件替换。它的基本用法有两种形式：

1. 基本形式：选择元素

np.where(condition, x, y)

示例：

import numpy as np

arr = np.array([1, 2, 3, 4, 5])

# 如果元素大于 3，替换为 10，否则替换为 0
result = np.where(arr > 3, 10, 0)
print(result)

输出：

[ 0  0  0 10 10]

np.where(condition)

示例：

arr = np.array([1, 2, 3, 4, 5])

# 查找大于 3 的元素的索引位置
indices = np.where(arr > 3)
print(indices)

输出：

(array([3, 4]),)

这里，array([3, 4]) 表示数组中值大于 3 的元素的索引是 3 和 4。

在多维数组中，np.where() 同样适用。它会分别返回每一维上满足条件的索引。

示例：

arr = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])

# 查找大于 5 的元素的索引位置
indices = np.where(arr > 5)
print(indices)

输出：

(array([1, 1, 2, 2]), array([2, 2, 1, 2]))

这里，返回的元组表示满足条件的位置，在第一个维度（行）上索引为 [1, 1, 2, 2]，在第二个维度（列）上索引为 [2, 2, 1, 2]。