Python基础——np.where

Python

按条件查找数列中满足要求的值和索引，如：

#-------------------------------------------------------------------------

# 产生一个乱序的随机序列

>>> import numpy as np

>>> import random

>>> A = np.arange(10)

>>> A

[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]

>>> random.shuffle(A)

>>> A

array([1, 7, 4, 9, 2, 3, 6, 0, 8, 5])

#---------------------------------------------------------------------------

>>> idx = [idx for (idx, val) in enumerata(A) if val > 5]

>>> idx

[1, 3, 6, 8]

>>> val = [val for (idx, val) in enumerata(A) if val > 5]

>>> val

[7, 9, 6, 8]

#----------------------------------------------------------------------------

np.where()

numpy中能够返回符合某一条件的下标函数：np.where()，只不过np.where()并不接受list类型的参数。

np.where()可用于接收3个参数，用于三目运算；也可用于接收1个参数，直接返回符合要求的下标。

#-----------------------------------------------------------------------------

>>> import numpy as np

>>> A

array([1, 7, 4, 9, 2, 3, 6, 0, 8, 5])

>>> idx = np.where(A > 5)

(array([1, 3, 6, 8], dtype=int64),)

>>> A[idx]              #不建议如此操作

array([7, 9, 6, 8])

>>> A[A > 5]           #推荐操作

array([7, 9, 6, 8])

#----------------------------------------------------------------------------

注意：这种情况下，np.where() 用于返回断言成立时的索引，返回值的形式为 arrays of tuple，由 np.array 构成的 tuple，一般 tuple 的 len 为2（当判断的对象是多维数组时），哪怕是一维数组返回的仍是 tuple，此时tuple 的 len 为 1。如：

—— np.where()[0] 表示行的索引；

—— np.where()[1] 则表示列的索引；

np.where()用于三目运算：

#---------------------------------------------------------------------------

#如果A%2==0成立，则执行A+1，否则执行A-1

>>> A

array([1, 7, 4, 9, 2, 3, 6, 0, 8, 5])

>>> B = np.where(A%2 == 0, A+1, A-1)   # 偶＋1，奇-1

>>> B

array([0, 6, 5, 8, 3, 2, 7, 1, 9, 4])

#---------------------------------------------------------------------------

参考文献：http://blog.csdn.net/lanchunhui/article/details/49489205

参考文献：https://docs.scipy.org/doc/numpy/reference/generated/numpy.where.html