【numpy】np.where()

np.where(condition [, x, y])

Return elements chosen from x or y depending on condition.
Parameters：
condition - array_like, bool Where True, yield x, otherwise yield y.
x, y - array_like Values from which to choose. x, y and condition need to be broadcastable to some shape.

Returns
out: ndarray An array with elements from x where condition is True, and elements from y elsewhere.

对于1-D array，相当于 [xv if c else yv for c, xv, yv in zip(condition, x, y)]

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

也可用于多维数组：

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

可以看出，多维数组也是元素级的判断。
该示例中，condition的shape 等于 x的shape 等于y的shape，三者对应位置为一组进行计算。

但是shape相同并不是必须的，当shape不同的时候，会进行broadcast。
官网上还有关于broadcast的示例。

# 示例1
>>>x, y = np.ogrid[:3, :4]
>>>np.where(x < y, x, 10 + y)  # both x and 10+y are broadcast
array([[10,  0,  0,  0],
       [10, 11,  1,  1],
       [10, 11, 12,  2]])

# 示例2

>>>a = np.array([[0, 1, 2],
              [0, 2, 4],
              [0, 3, 6]])
>>>np.where(a < 4, a, -1)  # -1 is broadcast
array([[ 0,  1,  2],
       [ 0,  2, -1],
       [ 0,  3, -1]])
       

 

当x,y缺省时

When only condition is provided, this function is a shorthand for np.asarray(condition).nonzero(). Using nonzero directly should be preferred, as it behaves correctly for subclasses. The rest of this documentation covers only the case where all three arguments are provided.

>>>aa = np.arange(8).reshape(2,2,2)
>>>aa
array([[[0, 1],
        [2, 3]],

       [[4, 5],
        [6, 7]]])
        
>>>np.where(aa>5)
(array([1, 1], dtype=int64),
 array([1, 1], dtype=int64),
 array([0, 1], dtype=int64))
# 输出的是满足条件的坐标
# 需要注意的是，这里的每一个array代表每一个维度。也就是说，坐标是要竖着看的。这里符合条件的坐标是：
#[1,1,0] 和 [1,1,1]

 
 

Reference
官网： https://numpy.org/doc/stable/reference/generated/numpy.where.html