本文转载自微信公众号「 Python作业辅导员」,作者天元浪子 。转载本文请联系 Python作业辅导员公众号。
在Python语言中,in是一个使用频率非常高的操作符,用于判断对象是否位于字符串、元组、列表、集合或字典中。in操作和人的思维方式高度吻合,写起来近乎于自然语言,充分体现了Python的哲学理念。
- >>> 'or' in 'hello world'
- True
- >>> 5 in {1,2,3,4}
- False
- >>> 'age' in {'name':'Mike', 'age':18}
- True
有趣的是,除了R、javascript、SQL外,包括C/C++在内的主流语言几乎都不支持in操作。这或许可以解释为什么Python语言被认为是最容易学习的编程语言。
习惯了使用Python的in操作符,有时就会自然而然地应用到NumPy数组操作上。比如,下面的写法看起来没有任何问题。
- >>> import numpy as np
- >>> a = np.arange(9)
- >>> a
- array([0, 1, 2, 3, 4, 5, 6, 7, 8])
- >>> 5 in a
- True
- >>> 10 in a
- False
不过,当我尝试在np.where()函数中使用in操作符的时候,出现了意外。
- >>> np.where(a>5)
- (array([6, 7, 8], dtype=int64),)
- >>> np.where(a%2==0)
- (array([0, 2, 4, 6, 8], dtype=int64),)
- >>> np.where(a in [2,3,5,7])
- Traceback (most recent call last):
- File "<pyshell#111>", line 1, in <module>
- np.where(a in [2,3,5,7])
- ValueError: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()
使用a>5或者a%2==0作为条件,np.where()函数没有问题,但是,使用a in [2,3,5,7],np.where()就会抛出异常。即便写成下面这样,也不能得到期望的结果。
- >>> np.where(a in np.array([2,3,5,7]))
- (array([], dtype=int64),)
难道NumPy不支持两个数组之间的in操作吗?不,强大到宇宙无敌的NumPy,怎么会不支持数组之间的in操作呢?NumPy不但支持,而且支持得很好。
- >>> p = np.array([2,3,5,7]) # 质数数组
- >>> np.in1d(a, p) # 返回a的每一个元素是否是质数的布尔数组
- array([False, False, True, True, False, True, False, True, False])
- >>> np.where(np.in1d(a, p)) # 返回数组a中质数的索引序号
- (array([2, 3, 5, 7], dtype=int64),)
- >>> np.where(np.in1d(a, p), -1, a) # 返回数组a中的质数全部替换为-1的结果
- array([ 0, 1, -1, -1, 4, -1, 6, -1, 8])
np.in1d()的参数如果是多维数组,将被自动扁平化,且返回的布尔数组也是扁平化的一维数组。
- >>> np.in1d(a.reshape((3,3)), p)
- array([False, False, True, True, False, True, False, True, False])
如果np.in1d()的参数是多维的,且期望返回和原数组结构相同的布尔数组,则应使用np.isin()函数。
- >>> np.isin(a.reshape((3,3)), p)
- array([[False, False, True],
- [ True, False, True],
- [False, True, False]])
- >>> np.where(np.isin(a.reshape((3,3)), p))
- (array([0, 1, 1, 2], dtype=int64), array([2, 0, 2, 1], dtype=int64))
若是期望得到两个数组的交集而不是交集元素的索引,下面两种方式都可行。
- >>> a[np.where(np.isin(a, p))]
- array([2, 3, 5, 7])
- >>> np.intersect1d(a, p)
- array([2, 3, 5, 7])
第二种方式直接使用np.intersect1d()函数得到两个数组的交集,且自动排序。不过,我更喜欢第一种方式。