numpy.in1d、numpy.where的使用

1. numpy.in1d

函数功能：检查第一个数组中的每个元素是否也存在于第二个数组中。
函数原型：numpy.in1d(ar1, ar2, assume_unique=False, invert=False)
函数返回值：返回与ar1长度相同的布尔数组，如果ar1的元素在ar2中，则返回True，否则返回False。
参数：invert ：bool, 默认值为False。 如果为True，则返回的数组中的值将被反转。即在ar1中的一个元素位于ar2中时为False，否则为True。

>>> test = np.array([0, 1, 2, 5, 0])
>>> states = [0, 2]
>>> mask = np.in1d(test, states)
>>> mask
array([ True, False,  True, False,  True])
>>> test[mask]
array([0, 2, 0])
>>> mask = np.in1d(test, states, invert=True)
>>> mask
array([False,  True, False,  True, False])
>>> test[mask]
array([1, 5])

2. numpy.where

函数功能：给出数组中满足给定条件的元素的索引
函数原型：numpy.where(condition[, x, y])

注意：此处 [ ] 中的内容为可选参数。即调用函数时，x和y是可选参数，condition(判断条件)是必写参数。因此 numpy.where 大致可以分为两种用法：np.where(condition)和np.where(condition, x, y)。
下面分别进行介绍它们的使用方法：

2.1 np.where(condition)

函数调用时只给出了condition，则输出满足条件 (即非0) 元素的坐标。坐标以tuple的形式给出，通常原数组有多少维，输出的tuple中就包含几个数组，分别对应符合条件元素的各维坐标。
下面给出几个例子来具体说明：
例1：给出一个一维数组的例子：

>>> a = np.arange(2,20,3)
>>> a
array([ 2,  5,  8, 11, 14, 17])

数组 a 中大于 7 的数包括 8、11、14、 17。它们在数组中对应的索引为 2，3，4，5。因此

>>> np.where(a > 7)
(array([2, 3, 4, 5], dtype=int64),)

例2：给出一个多维数组的例子：

>>> x = np.arange(9.).reshape(3, 3)
>>> x
array([[ 0.,  1.,  2.],
       [ 3.,  4.,  5.],
       [ 6.,  7.,  8.]])
>>> np.where( x > 5 )
(array([2, 2, 2], dtype=int64), array([0, 1, 2], dtype=int64))

数组 x 中大于 5 的数包括 6、7、8。它们在数组中对应的索引为 第2个元素的第1、2、3个位置。输出的元组中包含2个数组，每个数组代表对应维的坐标。对于二维数组就是，第一个数组代表行坐标信息，第二个数组代表列坐标信息。
元素 6、7、8对应的数组索引为 (2,0), (2,1), (2,2)。即

array([2, 2, 2], dtype=int64),
array([0, 1, 2], dtype=int64)

如果想要得到数组 x 中大于 5 的数，则可以用数组索引的方式，即

>>> x[np.where( x > 5 )]
array([ 6.,  7.,  8.])

如果要想得到的数组 shape 与原来数组 shape 相同，则可以使用如下方法

>>> np.where(x > 5, x, -1)
array([[-1., -1., -1.],
       [-1., -1., -1.],
       [ 6.,  7.,  8.]])

满足条件 x > 5 ，则输出 x 对应索引的元素，否则输出 -1。此次调用了 np.where(condition, x, y) 函数，下面进行具体介绍。

2.2 np.where(condition, x, y)

函数调用时只给出了condition，x 和 y 。满足条件(condition)，输出x，不满足输出y。

>>> np.where([[True, False], [True, True]],
...          [[1, 2], [3, 4]],
...          [[9, 8], [7, 6]])
array([[1, 8],
      [3, 4]])

其中，[[True, False], [True, True]] 是条件， [[1, 2], [3, 4]] 是 x, [[9, 8], [7, 6]] 是 y。True 即可视为满足条件。

条件中第0个元素中的第0个元素是true,那么取x中的相对应元素1;
条件中第0个元素中的第1个元素是false,那么取y中的相对应元素8;
条件中第1个元素中的第0个元素是ture,那么取x中相对应的元素3;
条件中第1个元素中的第1个元素是ture,那么取x中相对应的元素4;

所以最后的结果中取出的元素是1,8,3,4。

参考文献

1. https://www.cnblogs.com/massquantity/p/8908859.html
2. https://zhuanlan.zhihu.com/p/83208224