03.17 choose 函数实现条件筛选

choose 函数实现条件筛选

对于数组,我们有时候需要进行类似 switchcase 进行条件选择,此时使用 choose 函数十分方便:

1import numpy as np
1control = np.array([[1,0,1],
2                    [2,1,0],
3                    [1,2,2]])
4
5np.choose(control, [10, 11, 12])
array([[11, 10, 11],
       [12, 11, 10],
       [11, 12, 12]])

在上面的例子中,choose0,1,2 对应的值映射为了 10, 11, 12,这里的 0,1,2 表示对应的下标。

事实上, choose 不仅仅能接受下标参数,还可以接受下标所在的位置:

 1i0 = np.array([[0,1,2],
 2               [3,4,5],
 3               [6,7,8]])
 4i2 = np.array([[20,21,22],
 5               [23,24,25],
 6               [26,27,28]])
 7control = np.array([[1,0,1],
 8                    [2,1,0],
 9                    [1,2,2]])
10
11np.choose(control, [i0, 10, i2])
array([[10,  1, 10],
       [23, 10,  5],
       [10, 27, 28]])

这里,control 传入第一个 1 对应的是 10,传入的第一个 0 对应于 i0 相应位置的值即 1,剩下的以此类推。

下面的例子将数组中所有小于 10 的值变成了 10

1a = np.array([[ 0, 1, 2], 
2              [10,11,12], 
3              [20,21,22]])
4
5a < 10
array([[ True,  True,  True],
       [False, False, False],
       [False, False, False]], dtype=bool)
1np.choose(a < 10, (a, 10))
array([[10, 10, 10],
       [10, 11, 12],
       [20, 21, 22]])

下面的例子将数组中所有小于 10 的值变成了 10,大于 15 的值变成了 15。

1a = np.array([[ 0, 1, 2], 
2              [10,11,12], 
3              [20,21,22]])
4
5lt = a < 10
6gt = a > 15
7
8choice = lt + 2 * gt
9choice
array([[1, 1, 1],
       [0, 0, 0],
       [2, 2, 2]])
1np.choose(choice, (a, 10, 15))
array([[10, 10, 10],
       [10, 11, 12],
       [15, 15, 15]])