Pandas中条件滚动累加的向量化实现

import pandas as pd
import numpy as np
from tqdm import tqdm

# 示例数据
data = pd.DataFrame({
    'shop': [1, 2, 1, 2, 1, 2, 1, 2, 1, 2] * 3,
    'execution_date': pd.to_datetime(pd.date_range(start='2023-11-01', periods=30, freq='D')),
    'item': np.random.choice(['stuff', 'other', 'another'], size=30)
})
data.loc[data.index % 3 == 0, 'item'] = 'stuff'
data.loc[data.index % 7 == 0, 'item'] = 'stuff'

cum_items = []
data.execution_date = pd.to_datetime(data.execution_date, errors="coerce")
subset = data[["shop", "execution_date", "item"]]

# 预处理：将需要计数的条件转换为数值（1或0）
# 假设我们要计算 'item' 列中 'stuff' 的出现次数
subset['is_stuff'] = (subset['item'] == 'stuff').astype(int)

for i, row in tqdm(subset.iterrows(), total=len(subset)):
    # 过滤条件：同一商店，且日期严格早于当前日期减去15天
    fdf = (subset
           .loc[(subset.shop == row["shop"])]
           .loc[subset.execution_date < (row["execution_date"] - pd.Timedelta(15, unit='d'))]
          )

    if len(fdf) == 0: 
        cum_items.append(np.nan)
    else: 
        cum_items.append(fdf['is_stuff'].sum()) # 计算 'stuff' 的出现次数

data["cum_items_iter"] = cum_items
print("迭代计算结果 (部分):")
print(data[['shop', 'execution_date', 'item', 'cum_items_iter']].head(10))

import pandas as pd
import numpy as np

# 示例数据（与上文相同）
data = pd.DataFrame({
    'shop': [1, 2, 1, 2, 1, 2, 1, 2, 1, 2] * 3,
    'execution_date': pd.to_datetime(pd.date_range(start='2023-11-01', periods=30, freq='D')),
    'item': np.random.choice(['stuff', 'other', 'another'], size=30)
})
data.loc[data.index % 3 == 0, 'item'] = 'stuff'
data.loc[data.index % 7 == 0, 'item'] = 'stuff'

# 1. 确保日期列为datetime类型
data.execution_date = pd.to_datetime(data.execution_date, errors="coerce")

# 2. 预处理：将需要计数的条件转换为数值（1或0）
# 假设我们要计算 'item' 列中 'stuff' 的出现次数
data['is_stuff'] = (data['item'] == 'stuff').astype(int)

# 3. 排序：按商店和日期排序，这对于滚动操作至关重要
# 注意：groupby_rolling 可能会打乱原始顺序，因此在计算前排序并保留原始索引是好习惯
subset_sorted = data.sort_values(['shop', 'execution_date']).copy()

# 4. 使用 groupby_rolling 进行向量化计算
# '15D' 定义了一个15天的窗口，包括当前日期并向前追溯15天
# on='execution_date' 指定了滚动基于的日期列
data['cum_items_vec'] = (subset_sorted.groupby('shop')
                                      .rolling('15D', on='execution_date')['is_stuff']
                                      .sum()
                                      .reset_index(level=0, drop=True) # 移除groupby引入的shop级别索引
                                      .reindex(data.index) # 将结果重新对齐到原始DataFrame的索引
                                     )

# 由于滚动窗口默认包含当前点，如果需要排除当前点，可以进一步处理
# 例如，如果'is_stuff'在当前日期为1，则减去它
# data['cum_items_vec'] = data['cum_items_vec'] - data['is_stuff']

print("\n向量化计算结果 (部分):")
print(data[['shop', 'execution_date', 'item', 'is_stuff', 'cum_items_vec']].head(10))

# 比较迭代和向量化结果（如果窗口定义一致）
# 注意：本例中，迭代方法和向量化方法对时间窗口的定义略有不同
# 迭代方法：D' < D - 15天 （严格早于15天前）
# 向量化方法：[D - 15天, D] （过去15天，包含当前日期）
# 因此，结果通常不会完全相同，但向量化方法解决了“过去N天”的常见需求