Pandas 滚动条件函数

Question

我在熊猫中滚动使用.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 有关 谢谢

Answer 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

其中n 是窗口大小（在本例中为 3）。这将返回所需的确切输出。

最好的！

Answer 2

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

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

如果您在 any 功能方面要求更笼统 - 我的建议是 - 始终尝试矢量化，通常避免 .apply(...)

Answer 3

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

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 的值。