如何用Pandas按组计算日期与组内中心日期的天数差

问题说明 #

我有如下示例DataFrame：

import pandas as pd

data = [['A', '2022-09-01', 2], ['A', '2022-09-02', 1], ['A', '2022-09-04', 3], ['A', '2022-09-06', 2],
        ['A', '2022-09-07', 1], ['A', '2022-09-07', 3], ['A', '2022-09-08', 4], ['A', '2022-09-08', 2],
        ['B', '2022-09-01', 2], ['B', '2022-09-03', 4], ['B', '2022-09-04', 2], ['B', '2022-09-05', 2],
        ['B', '2022-09-07', 1], ['B', '2022-09-08', 1], ['B', '2022-09-10', 1]]
df = pd.DataFrame(data = data, columns = ['group', 'date', 'value'])

希望新增一列diff_days，用于展示每条记录的日期与其所在组的中心日期的天数差值，中心日期对应的diff_days值为0。中心日期通过公式「最小日期+(最大日期-最小日期)/2」计算得出，例如A组的中心日期为2022-09-04，B组为2022-09-05，期望输出如下：

data = [['A', '2022-09-01', 2, -3], ['A', '2022-09-02', 1, -2], ['A', '2022-09-04', 3, 0], ['A', '2022-09-06', 2, 2],
        ['A', '2022-09-07', 1, 3], ['A', '2022-09-07', 3, 3], ['A', '2022-09-08', 4, 4], ['A', '2022-09-08', 2, 4],
        ['B', '2022-09-01', 2, -4], ['B', '2022-09-03', 4, -2], ['B', '2022-09-04', 2, -1], ['B', '2022-09-05', 2, 0],
        ['B', '2022-09-07', 1, 2], ['B', '2022-09-08', 1, 3], ['B', '2022-09-10', 1, 5]]
df_desired = pd.DataFrame(data = data, columns = ['group', 'date', 'value', 'diff_days'])

解决方案 #

步骤1：转换日期列类型 #

首先需要将date列从字符串转换为Pandas的datetime类型，否则无法进行日期运算：

df['date'] = pd.to_datetime(df['date'])

步骤2：按组计算中心日期并求差值 #

使用groupby结合transform方法，为每个组计算中心日期，再与每条记录的日期求天数差：

df['diff_days'] = df.groupby('group')['date'].transform(
    lambda x: x - (x.min() + (x.max() - x.min()) / 2)
).dt.days

完整代码 #

import pandas as pd

data = [['A', '2022-09-01', 2], ['A', '2022-09-02', 1], ['A', '2022-09-04', 3], ['A', '2022-09-06', 2],
        ['A', '2022-09-07', 1], ['A', '2022-09-07', 3], ['A', '2022-09-08', 4], ['A', '2022-09-08', 2],
        ['B', '2022-09-01', 2], ['B', '2022-09-03', 4], ['B', '2022-09-04', 2], ['B', '2022-09-05', 2],
        ['B', '2022-09-07', 1], ['B', '2022-09-08', 1], ['B', '2022-09-10', 1]]
df = pd.DataFrame(data = data, columns = ['group', 'date', 'value'])

# 转换日期类型
df['date'] = pd.to_datetime(df['date'])

# 计算diff_days列
df['diff_days'] = df.groupby('group')['date'].transform(
    lambda x: x - (x.min() + (x.max() - x.min()) / 2)
).dt.days

print(df)

运行后得到的结果与期望输出完全一致。

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