Pandas按相近日期范围分组并保留原始DataFrame的方法

原始数据与需求 #

原始DataFrame #

import pandas as pd
df = pd.DataFrame([
#score    start         end        length
[1379, '2018-08-20', '2021-08-16',  156],
[1374, '2018-06-25', '2021-08-23',  165],
[1374, '2018-07-02', '2021-08-23',  164],
[1343, '2018-07-30', '2021-08-23',  160],
[1342, '2019-06-03', '2021-08-16',  115],
[1342, '2019-06-10', '2021-08-16',  114],
[1333, '2018-07-09', '2021-08-23',  163],
[1329, '2018-08-06', '2021-08-23',  159],
[1326, '2018-07-23', '2021-08-23',  161],
[1316, '2018-06-25', '2021-08-16',  164],
[1316, '2018-07-02', '2021-08-16',  163],
[1316, '2018-07-09', '2021-08-16',  162],
[1316, '2018-08-27', '2021-08-16',  155],
[1311, '2018-09-03', '2021-07-12',  149],
[1311, '2018-09-03', '2021-07-19',  150],
[1308, '2018-08-20', '2021-08-02',  154],
[1308, '2018-08-20', '2021-08-09',  155],
[1308, '2018-08-20', '2021-08-23',  157],
[1297, '2019-08-12', '2021-08-23',  106],
[1296, '2018-09-17', '2021-08-09',  151]
], columns=['score', 'start', 'end', 'length'])

需求 #

对start和end列中相差小于2个月的日期进行分组，获取每组的最大score，同时保留原始所有列（score、start、end、length）。

原方法的问题 #

你之前使用的pd.Grouper(freq='2M')分组存在两个核心问题：

• 会将日期自动对齐到2个月周期的末尾（例如原始的2018-08-20会被归为2018-08-31），丢失了原始的日期值
• 分组后仅提取了score的最大值，直接丢弃了length列以及原始日期列

解决方案 #

以下提供两种符合需求的可行方案：

方案1：按固定2个月周期分组，保留原始行 #

此方案延续你原有的分组逻辑（将start和end分别归入固定2个月周期），但保留所有原始列数据：

# 转换日期类型
df["start"] = pd.to_datetime(df["start"], format="%Y-%m-%d")
df["end"] = pd.to_datetime(df["end"], format="%Y-%m-%d")

# 生成分组键：将日期映射到所属2个月周期的起始点（避免对齐到月末）
df['start_group'] = df['start'].dt.to_period('2M').dt.start_time
df['end_group'] = df['end'].dt.to_period('2M').dt.start_time

# 方法1：仅保留每组中score最大的第一行
result = df.loc[df.groupby(['start_group', 'end_group'])['score'].idxmax()]

# 方法2：保留每组中所有score等于最大值的行
# max_scores = df.groupby(['start_group', 'end_group'])['score'].transform('max')
# result = df[df['score'] == max_scores]

# 删除临时分组列
result = result.drop(['start_group', 'end_group'], axis=1)

print(result)

说明

• 使用dt.to_period('2M').dt.start_time生成每个日期所属2个月周期的起始日期作为分组依据，既保证了分组逻辑和原方法一致，又不会修改原始日期值
• 通过idxmax()获取每组中score最大的行索引，再用loc提取完整的原始行数据
• 若同一组内有多个行的score均为最大值，可使用transform('max')的方式保留所有符合条件的行

方案2：自定义日期差聚类分组（非固定周期） #

如果你的需求是将start日期相差小于2个月，且end日期相差小于2个月的行归为同一组（而非按固定2个月周期划分），可以使用层次聚类实现：

from scipy.cluster.hierarchy import linkage, fcluster
import numpy as np

# 转换日期类型
df["start"] = pd.to_datetime(df["start"], format="%Y-%m-%d")
df["end"] = pd.to_datetime(df["end"], format="%Y-%m-%d")

# 将日期转换为时间戳数值（单位：纳秒）
start_ts = df['start'].view('int64')
end_ts = df['end'].view('int64')

# 定义2个月的时间差阈值（按30天估算，单位：纳秒）
two_month_threshold = 2 * 30 * 24 * 3600 * 10**9

# 层次聚类：单链接法，将时间差小于阈值的行聚为一类
Z = linkage(np.column_stack([start_ts, end_ts]), method='single')
df['cluster'] = fcluster(Z, t=two_month_threshold, criterion='distance')

# 提取每组中score最大的行
result = df.loc[df.groupby('cluster')['score'].idxmax()].drop('cluster', axis=1)

print(result)

说明

• 使用层次聚类的单链接法，自动将start和end差值均小于2个月的行归为同一组，更贴合“相差小于2个月”的字面需求
• 无需依赖固定周期，适合日期分布不规则的场景

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