Pandas滚动条件函数



我在Panda中滚动使用.apply.aggregate时遇到了一些问题(当然,假设这是解决问题的正确方法(。假设我有一个包含两列a和B的数据帧。如果a的值为1,我想创建一个包含B滚动平均值的列C。更一般地说,我希望能够在滚动的基础上应用自定义函数,其中一些条件涉及数据帧的几列(例如,当B>x和/或C=y等时,列a的滚动和(

import pandas as pd
import numpy as np
df2 = pd.DataFrame({'A':[1,1,1,0,0,0,1,1,1],'B': [50,40,50,-20,20,10,10,-5,-2]}, index = np.arange(9))

所需输出为(假设滚动窗口为3(:

df2 = pd.DataFrame({'A':[1,1,1,0,0,0,1,1,1],'B': [50,40,50,-20,20,10,10,-5,-2],
'C': [np.nan, np.nan, 46.67, 45, 50, np.nan, 10, 2.50, 1]}, index = np.arange(9))

我试图定义一个函数mean_1如下:

def mean_1(x):
return np.where(x['A'] == 1, np.mean(x['B']), np.nan)
df2['C'] = df2.rolling(3).apply(mean_1)

得到错误:'Series' object has no attribute 'A'我想这与文档中的raw = False有关感谢

您可以矢量化您的解决方案:

df2['C'] = df2['A'].eq(1).mul(df2['B']).rolling(3).sum()
.div(df2['A'].eq(1).rolling(3).sum())
.round(2)

如果你对任何函数的要求更一般——我的建议是——总是尝试向量化,通常避免.apply(...)

这里有一种接近所需输出的方法。

df2['C'] = df2.apply(lambda row: np.where(row['A']==1, row['B'], np.nan), axis=1).rolling(3, min_periods=1).apply(np.nanmean)

不同之处在于,上面给出了索引0和1的值。

您可以首先屏蔽'A'不为1的'B'值,然后应用滚动方法:

mask_map = df2.A != 1
df2['C'] = df2.B.mask(mask_map).rolling(3, min_periods=1).mean().round(2)

输出:

A   B      C
0  1  50  50.00
1  1  40  45.00
2  1  50  46.67
3  0 -20  45.00
4  0  20  50.00
5  0  10    NaN
6  1  10  10.00
7  1  -5   2.50
8  1  -2   1.00

请注意,第一个值不是NaN,因为我们指定了min_periods=1。这意味着我们取的是平均值,而不考虑缺失值的数量。所以,如果是这样的话,如果你真的想把第一个值设置为NaN,这可以用来完成

df2.iloc[:n-1, df2.columns.get_loc('C')] = np.nan

其中CCD_ 11是窗口大小(在这种情况下为3(。这将返回所需的精确输出。

最好!

最新更新