NumPy库学习之np.where函数

NumPy库学习之np.where函数

一、简介

np.where 是 NumPy 库中的一个函数，它与 Python 内置的 where 关键字不同。np.where 函数用于找出满足条件的元素的位置，并且可以执行条件选择，类似于 numpy.choose 或者 numpy.select。

二、语法和参数

语法：

np.where(condition, [x, y])

参数：

• condition: 一个布尔数组，np.where 会找出其中为 True 的元素的位置。
• x: 可选参数，当 condition 为 True 时，返回 x 中对应的元素。
• y: 可选参数，当 condition 为 False 时，返回 y 中对应的元素。如果不提供 y，则默认返回 condition 中 False 对应的位置。

返回值：

• 返回一个元组，包含两个数组：满足条件的元素的索引和不满足条件的元素的索引。

三、实例

3.1 找出数组中非零元素的位置

import numpy as np

# 创建一个数组
arr = np.array([1, 0, 2, 3, 0])

# 使用 np.where 查找非零元素的位置
indices = np.where(arr != 0)

print("Indices of non-zero elements:", indices)

输出：

Indices of non-zero elements: (array([0, 2, 3], dtype=int64),)

3.2 条件选择

# 使用 np.where 进行条件选择
choice = np.where(arr > 1, 10, -10)

print("Selected values:", choice)

输出：

Selected values: [-10 -10  10  10 -10]

3.3 多条件选择

import numpy as np

# 创建一个数组
arr = np.array([[1, 0, 2, 3, 0],
                [1, 0, 2, 3, 0],
                [1, 0, 2, 3, 0],
                [1, 0, 2, 3, 0],
                [1, 0, 2, 3, 0]])
conditions = [arr > 2, arr <= 2]

# 使用 np.where 进行多条件选择
x, y = np.where((arr >= 2) & (arr <= 3))

for i in range(len(x)):
    print(f'（{x[i]}, {y[i]}）')

输出：

（0, 2）
（0, 3）
（1, 2）
（1, 3）
（2, 2）
（2, 3）
（3, 2）
（3, 3）
（4, 2）
（4, 3）

四、注意事项

• np.where 返回的是满足条件的元素的索引，而不是元素本身。
• 如果只提供 condition 参数，返回的是布尔索引数组。
• 当使用 np.where 进行条件选择时，x 和 y 必须具有相同的形状或是标量。
• np.where 可以处理多维数组，但条件数组必须是一维的。
• 使用 np.where 可以进行高效的条件索引，但要注意数组的广播规则。