choose 函数实现条件筛选

对于数组，我们有时候需要进行类似 switch 和 case 进行条件选择，此时使用 choose 函数十分方便：

In [1]:

  1. import numpy as np

In [2]:

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

Out[2]:

  1. array([[11, 10, 11],
  2.        [12, 11, 10],
  3.        [11, 12, 12]])

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

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

In [3]:

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

Out[3]:

  1. array([[10,  1, 10],
  2.        [23, 10,  5],
  3.        [10, 27, 28]])

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

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

In [4]:

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

Out[4]:

  1. array([[ True,  True,  True],
  2.        [False, False, False],
  3.        [False, False, False]], dtype=bool)

In [5]:

  1. np.choose(a < 10, (a, 10))

Out[5]:

  1. array([[10, 10, 10],
  2.        [10, 11, 12],
  3.        [20, 21, 22]])

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

In [6]:

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

Out[6]:

  1. array([[1, 1, 1],
  2.        [0, 0, 0],
  3.        [2, 2, 2]])

In [7]:

  1. np.choose(choice, (a, 10, 15))

Out[7]:

  1. array([[10, 10, 10],
  2.        [10, 11, 12],
  3.        [15, 15, 15]])

原文: https://nbviewer.jupyter.org/github/lijin-THU/notes-python/blob/master/03-numpy/03.17-choose.ipynb