基于类别和滑动日期范围查找条件均值的最佳方法

Question

日期	类别	X值
...	红色	12
...	蓝色	19
...	...	...

假设你有一个像上面这样的Pandas数据框，怎么才能高效地计算每一行的X值总和，条件是这些行必须在同一个类别下，并且日期要在当前行日期的前后14天内？

我的数据集有50万行。

我尝试过的方法是对每一行都考虑整个数据集，这样在处理几千行以上的数据时就会很慢。

如果我考虑类别和日历月份，可以加快速度，但我想要的是一种滑动窗口的方式，随着每一行的变化而变化，而不是固定的日历月份。

def cohort_value(row):
    row_date = row['date']
    window_start_date = row_date - timedelta(days=14)
    window_end_date = row_date + timedelta(days=14)
    
    cohort_value = df[(df['date'] > start_date) & (df['date'] < end_date)]['value'].sum()
    
    return  cohort_value

df.apply(cohort_value, axis=1)

Answer 1

使用一个小的数据框，它的结构和示例一样：

data = {
    "Date": pd.date_range(start="2024-03-01", periods=5),
    "Category": ["Red", "Red", "Red", "Blue", "Red"],
    "X-Value": [12, 19, 15, 22, 18],
}

df = pd.DataFrame(data)

        Date Category  X-Value
0 2024-03-01      Red       12
1 2024-03-02      Red       19
2 2024-03-03      Red       15
3 2024-03-04     Blue       22
4 2024-03-05      Red       18

如果你把 Date 设置为索引，你就可以按 Category 分组，计算滚动窗口，并对每组的值进行求和。在下面的例子中，我使用的是3天的窗口（window="3D"），如果你想要每行前14天和后14天的数据，可以把这个值改成29。

df_rolling = (
    df.set_index("Date")
    .groupby("Category")["X-Value"]
    .rolling(window="3D", center=True)
    .sum()
    .reset_index(name="rolling_sum")
)

df = pd.merge(df, df_rolling, how="left")

        Date Category  X-Value  rolling_sum
0 2024-03-01      Red       12         31.0
1 2024-03-02      Red       19         46.0
2 2024-03-03      Red       15         34.0
3 2024-03-04     Blue       22         22.0
4 2024-03-05      Red       18         18.0