Please refresh the page if equations are not rendered correctly.
---------------------------------------------------------------

np.where()

Numpy.where（condition,x,y）当条件（condition）满足时为真（True），返回x。当条件（condition）不满足时为假（False），返回y。返回的结果也是一个ndarray数组

单个条件

根据条件改变原数组

import numpy as np

a = np.arange(12).reshape((3, 4))
[[ 0  1  2  3]
 [ 4  5  6  7]
 [ 8  9 10 11]]
###############################
print(np.where(a < 4, -1, 50))
array([[-1, -1, -1, -1],
       [50, 50, 50, 50],
       [50, 50, 50, 50]])
###############################
'''满足条件或不满足条件的元素的替换'''
np.where(a < 4, -1, a)
array([[-1, -1, -1, -1],
       [ 4,  5,  6,  7],
       [ 8,  9, 10, 11]])
print(np.where(a < 4, a, 50))
[[ 0  1  2  3]
 [50 50 50 50]
 [50 50 50 50]]
###############################
print(np.where(a < 4, True, False))
[[ True  True  True  True]
 [False False False False]
 [False False False False]]
type(np.where(a < 4, True, False))
numpy.ndarray
###############################
a < 4
array([[ True,  True,  True,  True],
       [False, False, False, False],
       [False, False, False, False]])

同时可以这么写：

a = np.arange(12).reshape((3, 4))
print(a)
[[ 0  1  2  3]
 [ 4  5  6  7]
 [ 8  9 10 11]]
a[a < 4] = -1
print(a)
[[-1 -1 -1 -1]
 [ 4  5  6  7]
 [ 8  9 10 11]]

np.where()返回一个新的ndarray数组，原数组不变。

import numpy as np

a = np.arange(12).reshape((3, 4))
print(a)
[[ 0  1  2  3]
 [ 4  5  6  7]
 [ 8  9 10 11]]

a_new = np.where(a < 4, -1, a)
print(a_new)
[[-1 -1 -1 -1]
 [ 4  5  6  7]
 [ 8  9 10 11]]

省略参数x,y时

缺失x和y参数的情况下，则输出满足条件（非0）元素的index，等价于 np.asarray(condition).nonzero()

import numpy as np

a = np.arange(12).reshape((3, 4))
print(np.where(a < 4))
# (array([0, 0, 0, 0], dtype=int64), array([0, 1, 2, 3], dtype=int64))
print(type(np.where(a < 4)))
#
print(list(zip(*np.where(a < 4))))
[(0, 0), (0, 1), (0, 2), (0, 3)] #True 的 index
a < 4
array([[ True,  True,  True,  True],
       [False, False, False, False],
       [False, False, False, False]])

多个条件

多个条件式时，用（）将其分开，条件式与条件式之间可以用 &，| 逻辑运算符进行连接。但不可以使用and，or等关键字链接。

import numpy as np

a = np.arange(12).reshape((3, 4))
array([[ 0,  1,  2,  3],
       [ 4,  5,  6,  7],
       [ 8,  9, 10, 11]])
######################################
print(np.where((a > 2) & (a < 6), -1, 100))
[[100 100 100  -1]
 [ -1  -1 100 100]
 [100 100 100 100]]
######################################
print(np.where((a > 2) & (a < 6) | (a == 7), -1, 100))
[[100 100 100  -1]
 [ -1  -1 100  -1]
 [100 100 100 100]]
######################################
print((a > 2) & (a < 6))
[[False False False  True]
 [ True  True False False]
 [False False False False]]
######################################
print((a > 2) & (a < 6) | (a == 7))
[[False False False  True]
 [ True  True False  True]
 [False False False False]]

元素的计算处理

np.where()也可以进行计算后，返回一个新的数组。

import numpy as np

a = np.arange(12).reshape((3, 4))
print(a)
[[ 0  1  2  3]
 [ 4  5  6  7]
 [ 8  9 10 11]]
print(np.where(a < 4, a * 100, a))
[[  0 100 200 300]
 [  4   5   6   7]
 [  8   9  10  11]]
print(np.where(a < 4, a + 100, a))
[[100 101 102 103]
 [  4   5   6   7]
 [  8   9  10  11]]