np.where函数

<1> np.where( condition )

输入只有条件数组condition的话，返回的是满足条件的元素索引。
1）condition为一维数组，返回值out为一个一维数组索引组成的元组；
2）condition为二维数组，返回值out为两个一维数组组成的元组，第一个数组为dim=0（行）的索引，第二个数组为dim=1（列）的索引。

condition数组： condition二维数组可以来自于数据数组的比较得到。

if __name__=="__main__":
    x = np.arange(0,20).reshape((4, 5))
    print(x)
    condition1 = x>4
    print(condition1)

#output：
[[ 0  1  2  3  4]
 [ 5  6  7  8  9]
 [10 11 12 13 14]
 [15 16 17 18 19]]
[[False False False False False]
 [ True  True  True  True  True]
 [ True  True  True  True  True]
 [ True  True  True  True  True]]

if __name__=="__main__":
    a = np.arange(0,5)
    print(a)
    aout = np.where(a>=2)
    print(aout, type(aout))
    print(aout[0])

    b = np.arange(0,20).reshape((4, 5))
    print(b)
    print(np.where(b>10))
    
#output:
[0 1 2 3 4]
(array([2, 3, 4]),) <class 'tuple'>
[2 3 4]
[[ 0  1  2  3  4]
 [ 5  6  7  8  9]
 [10 11 12 13 14]
 [15 16 17 18 19]]
(array([2, 2, 2, 2, 3, 3, 3, 3, 3]), array([1, 2, 3, 4, 0, 1, 2, 3, 4]))

用处： 通过np.where(condition)得到索引元组后可以直接用于获取对应索引的元素数值。

print(b[np.where(b>10)])
#output:
[11 12 13 14 15 16 17 18 19]

x = np.arange(12.).reshape(4, 3)
print(x)
positions = [2,6,10]
ins = np.isin(x,positions)
print(ins)
print(np.where(ins))

#output:
array([[  0.,   1.,   2.],
       [  3.,   4.,   5.],
       [  6.,   7.,   8.],
       [  9.,  10.,  11.]])
array([[False, False,  True],
       [False, False, False],
       [ True, False, False],
       [False,  True, False]], dtype=bool)
(array([0, 2, 3], dtype=int64), array([2, 0, 1], dtype=int64))

<2> np.where(condition, x, y)

输入为condition条件，以及与condition同shape的x与y。
对condition的每个元素True或者False，若为True，输出x中对应的元素，否则输出y中对应的元素。

x = np.arange(12.).reshape(4, 3)
y = np.inf
print(np.where(x>5, x, y))

#output:
[[inf inf inf]
 [inf inf inf]
 [ 6.  7.  8.]
 [ 9. 10. 11.]]