choose 函数实现条件筛选
对于数组,我们有时候需要进行类似 switch
和 case
进行条件选择,此时使用 choose 函数十分方便:
In [1]:
import numpy as np
In [2]:
control = np.array([[1,0,1],
[2,1,0],
[1,2,2]])
np.choose(control, [10, 11, 12])
Out[2]:
array([[11, 10, 11],
[12, 11, 10],
[11, 12, 12]])
在上面的例子中,choose
将 0,1,2
对应的值映射为了 10, 11, 12
,这里的 0,1,2
表示对应的下标。
事实上, choose
不仅仅能接受下标参数,还可以接受下标所在的位置:
In [3]:
i0 = np.array([[0,1,2],
[3,4,5],
[6,7,8]])
i2 = np.array([[20,21,22],
[23,24,25],
[26,27,28]])
control = np.array([[1,0,1],
[2,1,0],
[1,2,2]])
np.choose(control, [i0, 10, i2])
Out[3]:
array([[10, 1, 10],
[23, 10, 5],
[10, 27, 28]])
这里,control
传入第一个 1
对应的是 10,传入的第一个 0
对应于 i0
相应位置的值即 1
,剩下的以此类推。
下面的例子将数组中所有小于 10
的值变成了 10
。
In [4]:
a = np.array([[ 0, 1, 2],
[10,11,12],
[20,21,22]])
a < 10
Out[4]:
array([[ True, True, True],
[False, False, False],
[False, False, False]], dtype=bool)
In [5]:
np.choose(a < 10, (a, 10))
Out[5]:
array([[10, 10, 10],
[10, 11, 12],
[20, 21, 22]])
下面的例子将数组中所有小于 10 的值变成了 10,大于 15 的值变成了 15。
In [6]:
a = np.array([[ 0, 1, 2],
[10,11,12],
[20,21,22]])
lt = a < 10
gt = a > 15
choice = lt + 2 * gt
choice
Out[6]:
array([[1, 1, 1],
[0, 0, 0],
[2, 2, 2]])
In [7]:
np.choose(choice, (a, 10, 15))
Out[7]:
array([[10, 10, 10],
[10, 11, 12],
[15, 15, 15]])