np.where 是一个非常有用的 NumPy 函数,用于根据某个条件返回数组中的元素。它有两种主要的使用方式:
- 单条件选择:当只提供一个条件表达式时,np.where 返回满足条件的元素的索引。
- 双条件选择:当提供条件表达式以及两个数组(或标量值)时,np.where 在条件为 True 的位置返回第一个数组的对应元素,在条件为 False 的位置返回第二个数组的对应元素。
示例 1:单条件选择
import numpy as np
# 创建一个数组
arr = np.array([1, 2, 3, 4, 5])
# 找到所有大于 3 的元素的索引
indices = np.where(arr > 3)
print("Indices where condition is True:", indices)
输出将是:
Indices where condition is True: (array([3, 4]),)
示例 2:双条件选择
import numpy as np
# 创建一个数组
arr = np.array([1, 2, 3, 4, 5])
# 如果元素大于 3,替换为 10,否则替换为 -10
new_arr = np.where(arr > 3, 10, -10)
print("New array:", new_arr)
输出将是:
New array: [-10 -10 -10 10 10]
关键点:
- np.where 可以用于获取满足条件的元素索引或根据条件替换元素。
- 当用于元素替换时,np.where 可以接受数组或标量作为替换值。