where in 数组_一起学数据分析之NumPy(11)——将条件逻辑表述为数组运算

import numpy as np

numpy.where函数是三元表达式x if condition else y 的矢量化版本。假设我们有一个布尔数组和两个值数组：

In [5]:

xarr = np.array([1.1,1.2,1.3,1.4,1.5])yarr = np.array([2.1,2.2,2.3,2.4,2.5])cond = np.array([True,False,True,True,False])

假设我们想要根据cond中的值选取xarr和yarr中的值：当cond中的值为True时，选取xarr的值，否者选取yarr的值。列表推导式的写法应该是如下所示：

In [6]:

result = [(x if c else y) for x, y, c in zip(xarr, yarr, cond)]result

Out[6]:

[1.1, 2.2, 1.3, 1.4, 2.5]

这有几个问题。第一，它对大数组的处理速度不是很快(因为所有的工作都是有纯Python完成的)。第二，无法用于多维数组。若使用np.where,则可以将该功能写的非常简洁：

In [8]:

result = np.where(cond,xarr,yarr)result

Out[8]:

array([1.1, 2.2, 1.3, 1.4, 2.5])

np.where的第二个和第三个参数不必是数组，它们都是标量值。在数据分析工作中，where通常用于根据另一个数组而产生一个新的数组。假设有一个由随机数据组成的矩阵，你希望将所有正值替换为2，将所有负值替换为-2。若利用np.where,则会非常简洁：

In [12]:

from numpy.random import randnarr = randn(4,4)arr

Out[12]:

array([[-1.38896225,  1.26356721, -0.18527058, -0.01477591],       [ 0.50136218,  2.07348996,  0.45507673,  0.59779691],       [ 0.47809658, -1.19317622, -0.2513158 ,  1.0457735 ],       [-1.42690278, -0.82705128,  0.14392008, -0.51085531]])

In [13]:

np.where(arr>0,2,-2) #将正值设置为2，负值设置为-2

Out[13]:

array([[-2,  2, -2, -2],       [ 2,  2,  2,  2],       [ 2, -2, -2,  2],       [-2, -2,  2, -2]])

In [14]:

np.where(arr>0,2,arr) #将正值设置为2

Out[14]:

array([[-1.38896225,  2.        , -0.18527058, -0.01477591],       [ 2.        ,  2.        ,  2.        ,  2.        ],       [ 2.        , -1.19317622, -0.2513158 ,  2.        ],       [-1.42690278, -0.82705128,  2.        , -0.51085531]])

传递给where的数组大小可以不相等，甚至可以是标量值。

只有稍微动动脑子，你就能用where表述出更复杂的逻辑。想象一下这样的一个例子，我有两个布尔型数组cond1和cond2，希望根据4种不同的布尔值组合实现不同的赋值操作：

In [2]:

#纯Python的写法可能是这样的result = []cond1 = [True,False,True,False,True,False,False,True,True,False]cond2 = [True,True,False,False,False,True,True,False,False,False]for i in range(10):    if cond1[i] and cond2[i]:        result.append(0)    elif cond1[i]:        result.append(1)    elif cond2[i]:        result.append(2)    else:        result.append(3)result

Out[2]:

[0, 2, 1, 3, 1, 2, 2, 1, 1, 3]

虽然不是很明显，但是这个for循环确实可以被改写成一个嵌套的where表达式：

In [5]:

cond1 = np.array(cond1)cond2 = np.array(cond2)np.where(cond1&cond2, 0 ,np.where(cond1,1,np.where(cond2,2,3)))

Out[5]:

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