NumPy numpy.where()的使用

NumPy模块提供了一个函数numpy.where()，用于根据条件选择元素。根据条件，它返回从a或b中选择的元素。

例如，如果所有参数 -> 条件、a和b都传入numpy.where()，那么它将根据条件生成的布尔数组中的值，返回从a和b中选择的元素。

如果只提供条件，则此函数是函数np.asarray(condition).nonzero()的简写。尽管直接使用nonzero更好，因为它对于子类的行为是正确的。

语法

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

参数

以下是numpy.where()函数中的参数：

condition：array_like，bool类型

如果将此参数设置为True，则产生x，否则产生y。

x, y：array_like类型：

此参数定义要选择的值。x、y和condition需要能够进行广播到某种形状。

返回值

此函数返回一个数组，其中元素来自x中条件为True的位置，其他位置则来自y中的元素。

示例1：np.where()

import numpy as np
a=np.arange(12)
b=np.where(a<6,a,5*a)
b

在上面的代码中

• 我们使用别名np导入了numpy库。
• 我们使用np.arange()函数创建了一个数组’a’。
• 我们声明了变量’b’并将np.where()函数的返回值赋给了它。
• 我们将数组’a’传递给了函数。
• 最后，我们尝试打印变量’b’的值。

在输出中，满足条件的值从0到5保持不变，而其他值都乘以了5。

输出：

array([ 0,  1,  2,  3,  4,  5, 30, 35, 40, 45, 50, 55])

示例2：多维数组

import numpy as np
a=np.arange(12)
b=np.where([[True, False], [True, True]],[[1, 2], [3, 4]],[[9, 8], [7, 6]])
b

输出：

array([[1, 8],
        [3, 4]])

示例3：广播x、y和condition

import numpy as np
x, y = np.ogrid[:3, :4]
a=np.where(x > y, x, 10 + y)
a

输出：

array([[10, 11, 12, 13],
        [ 1, 11, 12, 13],
        [ 2,  2, 12, 13]])

在上面的代码中

• 我们使用别名np导入了numpy库。
• 我们使用np.arange()函数创建了一个数组’a’。
• 我们声明了变量’b’并将np.where()函数的返回值赋给了它。
• 我们将多维布尔数组作为条件传递，并将x和y作为整数数组传递。
• 最后，我们尝试打印变量b的值。

在输出结果中，如果满足条件，x的值将与y的值进行比较，然后打印x的值；否则，打印作为where()函数参数传递的y的值。

示例4：广播特定值

x=np.array([[0,1,2],[0,2,5],[0,4,8]])
y=np.where(x<4,x,-2)
y

输出：

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