在pandas数据帧中向量化条件赋值

Question

如果我有一个带有列x的数据帧df,并希望在伪代码中使用此值基于x的值创建列y

 if df['x'] <-2 then df['y'] = 1 
 else if df['x'] > 2 then df['y']= -1 
 else df['y'] = 0

我将如何实现这一目标.我认为np.where是最好的方法,但不知道如何正确编码.

Answer 1

一种简单的方法是首先分配默认值,然后执行2次loc调用:

In [66]:

df = pd.DataFrame({'x':[0,-3,5,-1,1]})
df
Out[66]:
   x
0  0
1 -3
2  5
3 -1
4  1

In [69]:

df['y'] = 0
df.loc[df['x'] < -2, 'y'] = 1
df.loc[df['x'] > 2, 'y'] = -1
df
Out[69]:
   x  y
0  0  0
1 -3  1
2  5 -1
3 -1  0
4  1  0

如果你想使用np.where那么你可以使用嵌套np.where:

In [77]:

df['y'] = np.where(df['x'] < -2 , 1, np.where(df['x'] > 2, -1, 0))
df
Out[77]:
   x  y
0  0  0
1 -3  1
2  5 -1
3 -1  0
4  1  0

所以这里我们定义第一个条件,其中x小于-2,返回1,然后我们有另一个np.where条件测试另一个条件,其中x大于2并返回-1,否则返回0

计时

In [79]:

%timeit df['y'] = np.where(df['x'] < -2 , 1, np.where(df['x'] > 2, -1, 0))

1000 loops, best of 3: 1.79 ms per loop

In [81]:

%%timeit
df['y'] = 0
df.loc[df['x'] < -2, 'y'] = 1
df.loc[df['x'] > 2, 'y'] = -1

100 loops, best of 3: 3.27 ms per loop

因此,对于此样本数据集,该np.where方法的速度是原来的两倍

Answer 2

用于np.select多种条件

np.select(condlist, choicelist, default=0)

• choicelist根据 中的相应条件返回 中的元素condlist。
• default当所有条件评估为 时，使用该元素False。

condlist = [
    df['x'] < -2,
    df['x'] > 2,
]
choicelist = [
    1,
    -1,
]
df['y'] = np.select(condlist, choicelist, default=0)

np.select比嵌套更具可读性np.where，但速度同样快：

^{df = pd.DataFrame({'x': np.random.randint(-5, 5, size=n)})}

Answer 3

pd.cut对于您定义范围并基于ranges您可以分配的范围来说，这是一个很好的用例labels：

df['y'] = pd.cut(df['x'], [-np.inf, -2, 2, np.inf], labels=[1, 0, -1], right=False)

输出

   x  y
0  0  0
1 -3  1
2  5 -1
3 -1  0
4  1  0