如何使用Pandas替代循环实现带日期过滤的分组中位数计算并合并DataFrame

原来的循环写法在数据量大时效率极低，因为每次循环都要对df2进行布尔索引筛选，时间复杂度是O(n²)，完全无法处理大规模数据。下面给出两种基于pandas向量化操作的优化方案，能把效率提升几个数量级：

方案一：合并+分组计算中位数（通用场景） #

这个方案逻辑直观，适合大多数情况，核心是通过合并、筛选、分组三步完成计算，全程避免循环：

步骤说明 #

1. 统一日期格式：把concert_date和date转换成datetime类型，并统一为当月最后一天（避免年月格式的日期默认取当月第一天导致的范围计算误差）。
2. 保留原始索引：给df1添加原始索引，方便后续把计算结果准确映射回去。
3. 合并数据：基于artist_id和region_id合并df1和df2，让每个df1的行与同组的所有df2行配对。
4. 筛选时间范围：只保留df2中日期在concert_date前3个月内（含当月）的行。
5. 分组计算中位数：按df1的原始索引分组，计算每组的popularity中位数。
6. 映射结果并补全缺失值：把中位数结果合并回df1，对没有匹配到数据的行填充默认值（比如示例中的44）。

代码实现 #

import pandas as pd

# 示例数据
df1 = pd.DataFrame({
    'artist_id': [12345, 33322],
    'concert_date': ['2019-10', '2018-11'],
    'region_id': [22, 44]
})

df2 = pd.DataFrame({
    'artist_id': [12345, 12345],
    'date': ['2019-10', '2019-11'],
    'region_id': [22, 44],
    'popularity': [76, 23]
})

# 1. 统一日期格式为当月最后一天
df1['concert_date'] = pd.to_datetime(df1['concert_date']) + pd.offsets.MonthEnd(0)
df2['date'] = pd.to_datetime(df2['date']) + pd.offsets.MonthEnd(0)

# 2. 保留原始索引
df1 = df1.reset_index(names='original_idx')

# 3. 合并df1和df2
merged = pd.merge(df1, df2, on=['artist_id', 'region_id'], how='left')

# 4. 筛选3个月内的日期范围
merged = merged.query('date >= concert_date - pd.DateOffset(months=3) and date <= concert_date')

# 5. 按原始索引分组计算中位数
median_results = merged.groupby('original_idx')['popularity'].median().reset_index()

# 6. 合并结果并补全缺失值
df1 = df1.merge(median_results, on='original_idx', how='left')
df1 = df1.rename(columns={'popularity': 'popularity_median_last3month'})
# 对无匹配数据的行填充默认值（示例中为44，可根据实际需求调整）
df1['popularity_median_last3month'] = df1['popularity_median_last3month'].fillna(44)

# 恢复原始结构
df1 = df1.drop('original_idx', axis=1)

print(df1)

方案二：滚动时间窗口（适合df2数据按时间连续的场景） #

如果df2中每个artist_id+region_id组的日期是连续的，可以用滚动时间窗口预计算中位数，再映射到df1，效率会更高：

代码实现 #

# 日期转换（同方案一）
df1['concert_date'] = pd.to_datetime(df1['concert_date']) + pd.offsets.MonthEnd(0)
df2['date'] = pd.to_datetime(df2['date']) + pd.offsets.MonthEnd(0)

# 对df2按分组排序，计算每个日期点的前3个月滚动中位数
df2_sorted = df2.sort_values(['artist_id', 'region_id', 'date'])
df2_rolling_median = df2_sorted.groupby(['artist_id', 'region_id']).rolling(
    window='3M', on='date'
)['popularity'].median().reset_index(name='popularity_median_last3month')

# 合并df1和预计算的中位数，匹配同分组且日期<=concert_date的最大日期对应的中位数
df1 = pd.merge_asof(
    df1.sort_values('concert_date'),
    df2_rolling_median.sort_values('date'),
    by=['artist_id', 'region_id'],
    left_on='concert_date',
    right_on='date',
    direction='backward'
)

# 补全缺失值
df1['popularity_median_last3month'] = df1['popularity_median_last3month'].fillna(44)
# 恢复原始顺序
df1 = df1.sort_index().drop('date', axis=1)

print(df1)

为什么优化后更快？ #

• 原来的循环是逐行遍历+多次全表筛选，时间复杂度为O(n²)；
• 优化后的方案用pandas内置的向量化操作（基于C语言实现），时间复杂度降到O(n log n)，数据量越大，效率提升越明显。

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