從Python到NumPy,細說最接近人類思維的in操作
本文轉載自微信公眾號「 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()函數得到兩個數組的交集,且自動排序。不過,我更喜歡第一種方式。