【numpy】numpy.where的使用

官方api介绍

几种用法

1. 仅输入condition

import numpy as np

a = np.array([0,5,4,1,9,7])
b = np.where(a>5)
print(b)

结果如下

它的结果是非0元素在原数组中的下标

2. 仅输入数组

1D数组

在官方note中，已经注明，如果是仅仅输入condition;那么它的作用等同于
​​​np.asarray(condition).nonzero()​​

如：

b = np.where([0,0,0,1,1,1,0,0,0])
print(b)

输出为：

2D数组

如果是2D，则返回为一个tuple，第一个值表示2D数组0维度的下标，第二个值表示2D数组1维度的下标

i = np.array([[False, False],
              [False, True],
              [True, True]])

print(np.where(i))

3. 三参数输入，并且每个参数都是1D

它等同于
​​​[xv if c else yv for c, xv, yv in zip(condition, x, y)]​​

如下面例子

import numpy as np

a = np.arange(9)
b = np.where(a<5, a, 0)
print(b)

4. 三参数输入，并且每个参数是2D

如果是2D，则标识对应的元素从对应的True、False数组中获取

例如，

b = np.where([[True, False], [True, True]],
         [[1, 2], [3, 4]],
         [[9, 8], [7, 6]])
print(b)

结果如下，

解释一下

2D数组

[True, False], [True, True]

2行2列来标识：

[ True, False

True, True]

它标识结果中的True从第二个参数数组中获取，False从第三个参数数组中获取。

总结

1个参数输出的是下标；3个参数输出的是对应的数值