详解Python中where()函数的用法

当你想在Python中快速筛选一些数据的时候，可以使用where()函数来帮助你实现这个任务，那么，下面我们来详解Python中where()函数的用法。

什么是where()函数

where()函数是NumPy中用于数组索引的函数，它返回索引数组，该索引数组的元素指示给定条件的元素的位置。where()函数的原型如下：

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

参数condition是需要满足的条件，可以是如下形式之一：

• 一个条件数组（数组的长度必须与其他参数一致）
• 一个条件表达式，该表达式将返回一个布尔类型的数组
• 一个只有一个维度的数组，其中所有非零元素都被视为True

参数x和y是可选的，它们是两个数组，用于根据condition的值从中选择元素。如果只提供一个参数，那么就会返回该参数中所有非零元素的索引。

where()函数的示例

下面我们来看两个简单的示例，这些示例将演示如何为一些数据选择元素，并使用where()函数来获取其索引。

示例1：从数组选择满足条件的元素

假设我们有一个数组arr，表示一些班级的成绩，我们想找到所有得分大于90分的学生及其得分。我们可以使用where()函数来获取这些学生的索引，如下所示：

import numpy as np

# 生成一个班级的成绩数组
arr = np.array([75, 92, 88, 65, 98, 84, 80, 91, 76, 89])

# 获取所有得分大于90分的学生及其得分
indices = np.where(arr > 90)
print("满足条件的学生的索引：", indices)
print("满足条件的学生的得分：", arr[indices])

运行这段代码，输出结果如下：

满足条件的学生的索引： (array([1, 4, 7]),)
满足条件的学生的得分： [92 98 91]

上述代码先生成一个班级的成绩数组，然后使用where()函数获取所有得分大于90分的学生的索引和对应得分。可以看到，where()函数返回了一个元组，这个元组只有一个元素，其值是一个一维数组，表示满足条件的学生的索引。arr[indices]语句返回满足条件的学生的得分。

示例2：从两个数组中选择符合条件的元素

假设我们有两个数组，分别表示一个人的年龄和体重，我们想查找所有年龄大于等于30岁且体重在150-160之间的人的信息，包括他们的年龄和体重。我们可以像下面这样使用where()函数来获取满足条件的人的索引：

import numpy as np

# 生成年龄和体重数组
ages = np.array([20, 25, 30, 35, 40])
weights = np.array([120, 140, 150, 160, 170])

# 获取年龄大于等于30岁且体重在150-160之间的人的信息
indices = np.where((ages >= 30) & (weights >= 150) & (weights <= 160))
print("满足条件的人的信息：", [(ages[i], weights[i]) for i in indices[0]])

运行这段代码，输出结果如下：

满足条件的人的信息： [(30, 150), (30, 160), (35, 150)]

上述代码先生成一个年龄和体重的数组，然后使用where()函数获取所有年龄大于等于30岁且体重在150-160之间的人的索引。然后依次遍历每个索引，输出满足条件的人的信息。

总结

通过上述示例，我们可以看到，在Python中使用where()函数可以轻松地选择满足条件的数组元素，并返回元素的索引。where()函数非常强大并且非常简单易用，在你需要快速筛选一些数据的时候可以考虑使用它。