如何向量化Pandas中按门店的15天前指定Item条件计数代码？

需求回顾 #

对DataFrame每一行，提取该行的日期与shop列值，筛选出同shop且日期早于当前行日期15天的所有行，统计其中item为'stuff'的出现次数。原代码使用iterrows()逐行循环，数据量较大时运行效率极低。

原代码 #

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

cum_items = []

data.execution_date = pd.to_datetime(data.execution_date, errors="coerce")

subset = data[["shop", "execution_date", "item"]]
for i, row in tqdm(subset.iterrows(), total=len(subset)):
    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.query("item == 'stuff'").shape[0])

data["cum_items"] = cum_items

向量化优化方案 #

核心思路是通过分组排序+向量化位置查找+累计求和替代逐行循环，时间复杂度从O(n²)降至O(n log n)，大幅提升运行效率：

1. 预处理数据：先标记出item为'stuff'的行，方便后续统计
2. 按店铺分组处理：对每个店铺的记录按日期排序，利用searchsorted快速定位符合日期条件的行范围，再通过累计求和得到目标次数

import pandas as pd
import numpy as np

# 1. 预处理：转换日期格式，创建商品标记列
data['execution_date'] = pd.to_datetime(data['execution_date'], errors="coerce")
data['is_stuff'] = np.where(data['item'] == 'stuff', 1, 0)

# 2. 定义分组处理函数
def compute_cumulative_stuff(group):
    # 按日期排序，确保时间顺序正确
    sorted_group = group.sort_values('execution_date').reset_index(drop=True)
    # 计算当前行日期往前推15天的截止日期
    cutoff_dates = sorted_group['execution_date'] - pd.Timedelta(days=15)
    # 用searchsorted找到每个截止日期在排序后日期列中的位置（左侧插入点）
    # 位置左侧的所有行都是早于截止日期的记录
    positions = sorted_group['execution_date'].searchsorted(cutoff_dates, side='left')
    # 计算标记列的累计和
    cum_sum = sorted_group['is_stuff'].cumsum()
    # 根据位置获取累计次数：位置为0时无符合条件的行，赋值NaN；否则取对应位置前的累计和
    sorted_group['cum_items'] = np.where(
        positions == 0, 
        np.nan, 
        cum_sum.take(positions - 1, fill_value=0)
    )
    # 恢复原数据的索引
    return sorted_group.set_index(group.index)

# 3. 应用到整个数据集
data = data.groupby('shop', group_keys=False).apply(compute_cumulative_stuff)
# 清理临时标记列
data = data.drop(columns='is_stuff')

优化说明 #

• 排序与searchsorted：每个店铺内按日期排序后，searchsorted能快速找到符合"早于当前行15天"的行边界，避免了逐行筛选的重复计算
• 累计求和：cumsum()是Pandas内置的高效向量化操作，一次计算就能得到所有行的累计统计结果
• 分组处理：groupby+apply确保每个店铺的统计独立进行，逻辑与原代码完全一致，但效率提升显著

内容的提问来源于stack exchange，提问作者Gustar8