Pandas如何向量化替代apply计算每行关联日期的count均值

现有如下格式的Pandas DataFrame：

import pandas as pd
df = pd.DataFrame(
{
"date":["2020/01/01","2020/01/02","2020/01/03","2020/01/04","2020/01/05"], 
"days_before": [["2019/12/31","2019/12/30"], ["2020/01/01", "2019/12/31"], ["2020/01/02","2020/01/01"], ["2020/01/03","2020/01/02"], ["2020/01/04","2020/01/03"]],
"count":[6,3,2,5,7]
})
print(df)

运行输出：

date               days_before  count
0  2020/01/01  [2019/12/31, 2019/12/30]      6
1  2020/01/02  [2020/01/01, 2019/12/31]      3
2  2020/01/03  [2020/01/02, 2020/01/01]      2
3  2020/01/04  [2020/01/03, 2020/01/02]      5
4  2020/01/05  [2020/01/04, 2020/01/03]      7

需求为新增before_avg列，列值取每行days_before字段包含的所有日期对应的count值的平均值。目前通过apply逐行遍历的方式可以实现，但需要性能更高的向量化实现方案，缩短执行耗时。
现有apply实现代码如下：

def calculateAvg(row):
   localDf = df[df["date"].isin(row["days_before"])]["count"]
   return localDf.mean()
df["before_avg"] = df.apply(lambda row: calculateAvg(row), axis=1)
print(df)

运行输出：

date               days_before  count  before_avg
0  2020/01/01  [2019/12/31, 2019/12/30]      6         NaN
1  2020/01/02  [2020/01/01, 2019/12/31]      3         6.0
2  2020/01/03  [2020/01/02, 2020/01/01]      2         4.5
3  2020/01/04  [2020/01/03, 2020/01/02]      5         2.5
4  2020/01/05  [2020/01/04, 2020/01/03]      7         3.5

核心逻辑是避免Python层逐行循环，全部使用Pandas底层C实现的批量操作完成计算，数据量越大性能优势越明显，十万级以上数据比apply快10~100倍。

通用实现（适配任意长度的days_before列表） #

# 构建日期到count值的哈希映射，查询复杂度O(1)
date_count_map = df.set_index("date")["count"].to_dict()

# 将列表格式的days_before拆分为多行，保留原行索引
exploded_days = df["days_before"].explode()

# 批量映射所有日期对应的count值，表中不存在的日期自动返回NaN
matched_count = exploded_days.map(date_count_map)

# 按原行索引分组求均值，直接赋值为新列，计算时自动忽略NaN
df["before_avg"] = matched_count.groupby(level=0).mean()

计算结果和原apply逻辑完全一致。

定长列表场景优化版 #

如果每行days_before的列表长度固定（比如示例中均为2个日期），可以直接将列表拆分为多列后批量计算，性能还能进一步提升：

date_count_map = df.set_index("date")["count"].to_dict()
# 直接把列表拆成多列
days_df = pd.DataFrame(df["days_before"].tolist(), index=df.index)
# 批量替换为count值后按行求均值
df["before_avg"] = days_df.replace(date_count_map).astype(float).mean(axis=1)

注意：两种实现都和原apply逻辑保持一致，表中不存在的日期会被自动忽略，不参与均值计算。

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