np.where

numpy.where

numpy.where([condition,x,y])
# 返回值out

根据condition选择x或y中的元素
 

Parameters

• condition：array_like, bool
  • 当condition为True时，选择x中的元素，反之y
• x，y：array_like
  • x，y，condition必须是broadcastable to some shape
• out：ndarray
  • 数组，yield x if condition==True else y
     

Notes

• 若所有参数都是1-D，函数等价于

[xv if c else yv for c, xv, yv in zip(condition,x,y)]

• 若仅提供了condition，函数等价于

condition.nonzeros()

 

Examples

>>>a=np.arange(12)
>>>a
array([ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11])
>>>np.where(a<5,a,a*10)
array([  0,   1,   2,   3,   4,  50,  60,  70,  80,  90, 100, 110])

可拓展到高维数组：

>>>np.where([[True,False],[True,True]],[[1,3],[5,7]],[[2,6],[6,8]])
array([[1, 6],
       [5, 7]])

x,y,condition的形状都是broadcast的

>>>a=np.array([[1,2,3],[4,5,6],[7,8,9]])
>>>np.where(a<4,a,-1)
array([[ 1,  2,  3],
       [-1, -1, -1],
       [-1, -1, -1]])