numpy.where()的使用方法_numpy where

numpy.where()主要用于条件筛选，即选择满足某些条件的行、列或元素。

有三种使用方法

1. 单参数numpy.where(condition)

numpy.where(condition) 当 where 内只有一个参数时，那个参数表示条件，当条件成立时，where 返回的是每个符合 condition 条件元素的坐标，返回的是以元组的形式

>>> a = np.array([2,4,6,8,10])
#只有一个参数表示条件的时候
>>> np.where(a > 5)
 
Out[5]: (array([2, 3, 4], dtype=int64),)

输出满足条件 (即非0) 元素的坐标 (等价于numpy.nonzero)。这里的坐标以tuple的形式给出，通常原数组有多少维，输出的tuple中就包含几个数组，分别对应符合条件元素的各维坐标。 

>>> a = np.arange(27).reshape(3,3,3)
>>> a
array([[[ 0,  1,  2],
        [ 3,  4,  5],
        [ 6,  7,  8]],
 
       [[ 9, 10, 11],
        [12, 13, 14],
        [15, 16, 17]],
 
       [[18, 19, 20],
        [21, 22, 23],
        [24, 25, 26]]])
 
>>> np.where(a > 5)
(array([0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2, 2, 2, 2]),
 array([2, 2, 2, 0, 0, 0, 1, 1, 1, 2, 2, 2, 0, 0, 0, 1, 1, 1, 2, 2, 2]),
 array([0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2]))
 
 
# 符合条件的元素为
	   [ 6,  7,  8]],
 
      [[ 9, 10, 11],
       [12, 13, 14],
       [15, 16, 17]],
 
      [[18, 19, 20],
       [21, 22, 23],
       [24, 25, 26]]]

2. 三参数numpy.where(condition, x, y)

np.where(condition, x, y) 当 where 内有三个参数时，第一个参数表示条件，当条件成立时 where 方法返回 x，当条件不成立时 where 返回 y

>>> import numpy as np
 
>>> x = np.random.randn(4, 4)
Out[4]: 
array([[-1.07932656,  0.48820426,  0.27014549,  0.43823363],
       [ 0.28400645,  0.89720027,  1.6945324 , -1.41739129],
       [ 0.55640566, -0.99401836, -1.58491355,  0.90241023],
       [-0.14711914, -1.21307824, -0.0509225 ,  1.39018565]])
>>> np.where(x > 0, 2, -2)
Out[5]: 
array([[-2,  2,  2,  2],
       [ 2,  2,  2, -2],
       [ 2, -2, -2,  2],
       [-2, -2, -2,  2]])

3. 多条件

多条件时，&表示与，|表示或。

例如a = np.where((0<a)&(a<5), x, y)，当0<a与a<5满足时，返回x的值，当0<a与a<5不满足时，返回y的值。注意x，y必须和a保持相同尺寸。此外多条件的括号必须带上。在这个例子中条件必须写成(0<a)&(a<5)，写在0<a & a<5是错误的。

# 满足 (data>= 0) & (data<=2)，则返回np.ones_like(data)，否则返回 np.zeros_like(data)
>>> import numpy as np
 
>>> data = np.array([[0, 2, 0], 
				     [3, 1, 2],
                     [0, 4, 0]])
>>> new_data = np.where((data>= 0) & (data<=2), np.ones_like(data), np.zeros_like(data))
>>> print(new_data)
Out[6]:
[[1 1 1]
 [0 1 1]
 [1 0 1]]

参考文献

python中np.where()的使用方法_赵孝正的博客-CSDN博客

numpy.where() 用法详解 - massquantity - 博客