分组内每行前后动态天数范围内唯一位置数量统计需求

给定包含GROUP（分组ID）、DATE（日期）、LOCATION（位置）的数据集，需要针对每个分组中的每一行，统计**过去n1天（含当日）和未来n2天（含当日）**内的唯一位置数量。要求支持动态设置n1和n2，且因数据集规模较大，必须兼顾性能。示例中n1=2、n2=3。

输入数据集 #

import pandas as pd

data = {
        'GROUP': ['A', 'A', 'A', 'A', 'A', 'A', 'A', 'A', 'A', 'B', 'B', 'B', 'B', 'B', 'B', 'B', 'B'],
        'DATE': ['1/01/2023', '1/01/2023', '1/01/2023', '2/01/2023', '2/01/2023', '3/01/2023', '3/01/2023', '3/01/2023', '3/01/2023', '1/06/2023', '2/06/2023', '3/06/2023', '4/06/2023', '5/06/2023', '5/06/2023', '5/06/2023', '6/06/2023'],
        'LOCATION': ['A1', 'A1', 'A2', 'A1', 'A2', 'A2', 'A3', 'A3', 'A4', 'B1', 'B2', 'B2', 'B3', 'B4', 'B2', 'B1', 'B1']
    }
df = pd.DataFrame(data)

期望输出 #

GROUP       DATE       LOCATION  LOCS_LAST_n1_DAYS  LOCS_NEXT_n2_DAYS
0      A        1/01/2023       A1                 1                 4
1      A        1/01/2023       A1                 1                 4
2      A        1/01/2023       A2                 2                 4
3      A        2/01/2023       A1                 2                 4
4      A        2/01/2023       A2                 2                 3
5      A        3/01/2023       A2                 2                 3
6      A        3/01/2023       A3                 3                 2
7      A        3/01/2023       A3                 3                 2
8      A        3/01/2023       A4                 4                 1
9      B        1/06/2023       B1                 1                 2
10     B        2/06/2023       B2                 2                 2
11     B        3/06/2023       B2                 1                 3
12     B        4/06/2023       B3                 2                 4
13     B        5/06/2023       B4                 2                 3
14     B        5/06/2023       B2                 3                 2
15     B        5/06/2023       B1                 4                 1
16     B        6/06/2023       B1                 3                 1

解决方案 #

步骤1：标准化日期格式并排序 #

先将DATE列转为datetime类型，同时按分组和日期排序，确保时间窗口计算的准确性：

df['DATE'] = pd.to_datetime(df['DATE'], format='%d/%m/%Y')
df = df.sort_values(['GROUP', 'DATE']).reset_index(drop=True)

步骤2：按分组聚合日期与位置数据 #

提前按分组聚合日期列表和对应位置的集合，避免逐行重复计算，提升性能：

grouped = df.groupby('GROUP').agg(
    dates=('DATE', list),
    loc_sets=('LOCATION', lambda x: [set([val]) for val in x])
).reset_index()

步骤3：编写窗口统计函数 #

定义函数处理单个分组的日期和位置数据，计算每个行对应的过去/未来窗口内的唯一位置数：

def calculate_window_locs(group_dates, group_loc_sets, n1, n2):
    result_last = []
    result_next = []
    date_series = pd.Series(group_dates)
    
    for idx, current_date in enumerate(group_dates):
        # 计算过去n1天的唯一位置数
        mask_last = (date_series >= current_date - pd.Timedelta(days=n1-1)) & (date_series <= current_date)
        locs_last = set().union(*[group_loc_sets[i] for i in mask_last[mask_last].index])
        result_last.append(len(locs_last))
        
        # 计算未来n2天的唯一位置数
        mask_next = (date_series >= current_date) & (date_series <= current_date + pd.Timedelta(days=n2-1))
        locs_next = set().union(*[group_loc_sets[i] for i in mask_next[mask_next].index])
        result_next.append(len(locs_next))
    
    return result_last, result_next

步骤4：应用函数并合并结果 #

设置n1、n2参数，遍历分组计算后将结果合并回原DataFrame：

n1 = 2
n2 = 3

last_locs = []
next_locs = []

for _, row in grouped.iterrows():
    l, n = calculate_window_locs(row['dates'], row['loc_sets'], n1, n2)
    last_locs.extend(l)
    next_locs.extend(n)

df['LOCS_LAST_n1_DAYS'] = last_locs
df['LOCS_NEXT_n2_DAYS'] = next_locs

# 可选：恢复原日期格式显示
df['DATE'] = df['DATE'].dt.strftime('%d/%m/%Y')
print(df)

性能优化说明 #

• 分组聚合预处理：减少全局遍历次数，将计算范围限定在分组内
• 集合去重合并：利用集合的高效去重特性，替代逐行统计唯一值
• 向量化日期筛选：使用pandas Series的掩码筛选，比循环判断效率更高

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