NumPy where()

import numpy as np

originalArray = np.array([1, -2, -3, 4, 5])

condition = originalArray < 0


# For each element of originalArray,
# if condition is True, use 0 as element in the resultant array
# if condition is False, use the corresponding element in the resultant array
result = np.where(condition, 0, originalArray )

print(result)

# Output :  [1 0 0 4 5]

where() 语法

where() 的语法是

numpy.where(condition, x, y)

where() 参数

where() 方法接受三个参数

• condition - 一个布尔值或一个数组
• x - 如果 condition 为 True 时取的值
• y - 如果 condition 为 False 时取的值

注意: 我们也可以向 np.where() 传递一个参数。要了解它，请访问下面的 np.where() 单参数用法 部分。

where() 返回值

where() 方法返回一个新的 NumPy 数组。

示例 1: numpy.where() 与两个数组

import numpy as np

x = np.array([1, 2, 3, 4])
y = np.array([10, 20, 30, 40])

test_condition = x < 3 

# if test_condition is True, select element of x
# if test_condition is False, select element of y
result = np.where(test_condition, x, y)

print(result)

输出

[1 2 30 40]

示例 2: numpy.where() 与运算

我们也可以使用 numpy.where() 对数组元素执行操作。

import numpy as np

x = np.array([-1, 2, -3, 4])

# test condition
test_condition = x > 0

# if test_condition is True, select element of x
# if test_condition is False, select  x * -1
result = np.where(test_condition, x, x * -1)

print(result)

输出

[1 2 3 4]

示例 3: where() 与数组条件

我们可以在 where() 方法中使用 array_like 对象（如列表、数组等）作为条件。

import numpy as np

x = np.array([[1, 2], [3, 4]])
y = np.array([[-1, -2], [-3, -4]])

# returns element of x when True 
# returns element of y when False
result = np.where([[True, True], [False, False]], x, y)

print(result)

# returns element of x when True 
# returns element of y when False
result = np.where([[True, False], [False, True]], x, y)


print(result)

输出

[[1 2]
 [-3 -4]]

[[1 -2]
 [-3 4]]

示例 4: where() 与多个条件

where() 方法中的测试条件可能包含多个条件。

我们使用

• | 运算符对多个条件执行 OR 操作
• & 运算符对多个条件执行 AND 操作

import numpy as np

x = np.array([1, 2, 3, 4, 5, 6, 7])

# if element is less than 2 or greater than 6, test condition is True
test_condition1 = (x < 2) | (x > 6)
# select element of x if test condition is True
# select 0 if test condition is False
result1 = np.where(test_condition1, x, 0)

print(result1)

# if element is  greater than 2 and less than 6, test condition is True
test_condition2 = (x > 2) & (x < 6)

# select element of x if test condition is True
# select 0 if test condition is False
result2 = np.where(test_condition2, x, 0)

print(result2)

输出

[1 0 0 0 0 0 7]
[0 0 3 4 5 0 0]

示例 5: where() 仅一个参数

如果向 numpy.where() 传递一个参数（测试条件），它会通过返回索引来告诉我们在给定数组中满足给定条件的那些位置。

import numpy as np

originalArray = np.array([0, 10, 20, 30, 40, 50, 60, 70])

# returns index of elements for which the test condition is True
result = np.where(originalArray > 30)

print(result)

输出

(array([4, 5, 6, 7]),)