np.where()的简单理解与用法

https://numpy.org/doc/stable/reference/generated/numpy.where.html

用法一

import numpy as np
np.where(condition, x, y)

condition中需要包含一个待处理的ndarray(这里记为A),那么对于A中的每个元素,如果满足条件,则将这个元素替换为x,否则,将这个元素替换为y。一个例子如下:

import numpy as np

A = np.array([1, 2, 3, 4, 5])
res = np.where(A > 3, 1, 0)
print(res)

输出:

[0 0 0 1 1]

即,对于A中的每个元素,逐个对比其是否大于3,是则将对应位置"替换"为1,不是则为0。


用法二

import numpy as np
np.where(condition)

类似的,对于A中的每个元素,检查其是否满足condition,如果是则返回其坐标。一个例子如下:

import numpy as np

A = np.array([1, 2, 3, 4, 5])
res = np.where(A > 3)
print(res)

输出:

(array([3, 4], dtype=int64),)

即下标3和下标4处的值满足条件。可以发现,这里返回的其实是一个tuple,这个tuple的维度与数组本身维度一致,用于返回多维的坐标。此外,通过简单套娃也可以将满足条件的值给取出来:

A = np.array([1, 2, 3, 4, 5])
res = np.where(A > 3)
print(A[res])

输出:

[4 5]

此时效果等同于直接使用A[A>3]

猜你喜欢

转载自blog.csdn.net/qq_40714949/article/details/124024679