Pandas滚动条件函数

Question

我在滚动使用 pandas 中的 .apply 或 .aggregate 时遇到了一些麻烦（当然假设这是解决我的问题的正确方法）。假设我有一个包含两列 A 和 B 的数据框。我想创建一个列 C，如果 A 的值为 1，它将包含 B 的滚动平均值。更一般地说，我希望能够应用滚动自定义函数，某些条件涉及数据框的多列（例如，当 B > x and/or 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))

期望的输出是（假设滚动 window 为 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

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

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

Answer 2

您可以先屏蔽 'B' 值，其中 'A' 不是 1，然后应用滚动方法：

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 是 window 大小（在本例中为 3）。这将 return 确切的期望输出。

最好！

Answer 3

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

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

如果您在 any 函数方面询问更一般的问题 - 我的建议是 - 始终尝试矢量化，通常避免 .apply(...)

Pandas滚动条件函数

Pandas Rolling Conditional Function

pandas

rolling-computation