重温python之np.where

np.where 是 NumPy 库中一个功能强大的函数，它允许用户根据指定的条件来选择数组中的元素或获取满足条件的元素的索引。以下是 np.where 的详细介绍：

功能与用途

np.where 主要用于以下两种情境：

1. 基于条件选择元素：根据给定的条件，从两个数组中选择元素来创建一个新的数组。如果条件为真，则选择第一个数组中的对应元素；如果条件为假，则选择第二个数组中的对应元素。
2. 获取满足条件的元素的索引：当只提供一个条件参数时，np.where 会返回满足该条件的元素的索引。

语法

np.where 的基本语法如下：

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

• condition：用于测试的元素。可以是布尔数组、布尔条件表达式或能够返回布尔数组的函数。
• x：当 condition 为真时选择的元素或数组。
• y：当 condition 为假时选择的元素或数组。

返回值

• 如果提供了 x 和 y 参数，则返回一个与 condition 形状相同的数组，其中包含了根据 condition 从 x 和 y 中选择的元素。
• 如果只提供了 condition 参数，则返回一个包含满足 condition 的元素索引的元组。元组的形状取决于输入数组的形状。

示例

示例 1：基于条件选择元素

import numpy as np

# 创建两个数组
x = np.array([1, 2, 3, 4, 5])
y = np.array([10, 20, 30, 40, 50])

# 定义条件
condition = x > 3

# 使用 np.where 根据条件从 x 和 y 中选择元素
result = np.where(condition, x, y)

print(result)  # 输出：[10 20 30  4  5]

在这个例子中，当 x 中的元素大于 3 时，result 中的对应位置是 x 的元素；否则是 y 的元素。

示例 2：获取满足条件的元素的索引

import numpy as np

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

# 定义条件
condition = arr > 3

# 使用 np.where 获取满足条件的元素的索引
indexes = np.where(condition)

print(indexes)  # 输出：(array([3, 4]),)

在这个例子中，indexes 是一个元组，其中包含了一个数组，该数组包含了所有满足 arr > 3 条件的元素的索引。

注意事项

• 当使用 np.where 根据条件从两个数组中选择元素时，x 和 y 必须具有与 condition 相同的形状，或者至少能够广播到 condition 的形状。
• 当只提供 condition 参数时，返回的索引元组的形状与输入数组的形状相匹配。例如，如果输入是一个二维数组，那么返回的索引元组将包含两个数组，分别对应行索引和列索引。

np.where 是一个强大且灵活的工具，能够处理各种条件选择和数据索引的任务，是 NumPy 库中非常有用的函数之一。