在熊猫数据框中矢量化条件赋值

Posted 2023-02-23

【中文标题】在熊猫数据框中矢量化条件赋值

【英文标题】：vectorize conditional assignment in pandas dataframe

【发布时间】：2015-05-07 22:32:10

【问题描述】：

如果我有一个带有 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 是最好的方法，但不确定如何正确编码。

【问题讨论】：

注意，在 pandas 16.0 中将有一种额外的方法来执行此操作（现在到期？），类似于 dplyr mutate：pandas-docs.github.io/pandas-docs-travis/…

查看this answer 使用np.where 有两个选择，np. select 有更多选择。

【参考方案1】：

对多个条件使用np.select

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

根据condlist中的对应条件返回choicelist中的元素。 当所有条件的计算结果为 False 时，使用 default 元素。

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))}

【参考方案2】：

这是 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

【讨论】：

我喜欢这个答案，但它比 np.where 选项慢，基于我机器上的 %timeit

【参考方案3】：

一种简单的方法是先分配默认值，然后执行 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 方法的速度是原来的两倍