Python之numpy.where()的用法

numpy.where()有两种用法

• 第一种用法:np.where(condition, x, y)
  满足条件(condition),输出x, 不满足输出y

import numpy as np
aa = np.arange(10)
np.where(aa, 1, -1)  # 0为False, 所以第一个输出-1

output:
array([-1, 1, 1, 1, 1, 1, 1, 1, 1, 1])

• 第二种用法:np.where(condition)
  只有条件(condition), 没有x和y,则输出满足条件(即非0)元素的坐标(等价于numpy.nonzero). 这里的坐标以tuple的形式给出,通常原数组有多少维,输出的tuple中就包含几个数组,分别对应符合条件元素的各维坐标.

import numpy as np
a = np.arange(27)
print(a)

t1 = np.where(a>5)
print(t1)
print(type(t1))
t2 = np.where(a>5)[0]
print(t2)
print(type(t2))

output:
[ 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]
(array([ 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22,
23, 24, 25, 26], dtype=int64),)
<class 'tuple'>
[ 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26]
<class 'numpy.ndarray'>