如何筛选从上个月起满足指定连续观测数的DataFrame unique_id

原始数据 #

import pandas as pd
from datetime import datetime

data = {
    'ds': ['2024-02-01', '2024-01-01', '2023-12-01', '2024-02-01', '2023-12-01'],
    'y': [500, 600, 700, 800, 500],
    'unique_id': [1, 1, 1, 2, 2]
}

input_df = pd.DataFrame(data)

需求说明 #

编写函数，根据参数min_consecutive_observations筛选出从上个月起倒推满足最小连续观测数要求的unique_id。例如当min_consecutive_observations == 2时，预期输出如下：

expected_output = pd.DataFrame({
    'ds': ['2024-02-01', '2024-01-01', '2023-12-01'],
    'y': [500, 600, 700],
    'unique_id': [1, 1, 1]
})

注：原预期输出中unque_id为笔误，已修正为unique_id

现有代码（已获取上个月日期） #

last_month_input = datetime.now().replace(day=1, hour=0, minute=0, second=0, microsecond=0) - pd.DateOffset(months=1)

解决方案 #

核心思路 #

1. 将日期列转换为datetime类型，确保日期计算准确
2. 按unique_id分组，对每个组的日期排序并去重（避免同一月份多条数据干扰连续性判断）
3. 计算相邻日期的月份差，验证是否连续
4. 检查从上个月开始倒推，是否存在连续min_consecutive_observations个月份的观测记录
5. 筛选符合条件的unique_id，返回对应的数据行

完整函数代码 #

def filter_consecutive_ids(input_df, min_consecutive_observations):
    import pandas as pd
    from datetime import datetime
    
    # 复制数据避免修改原表
    df = input_df.copy()
    # 转换日期列为datetime类型
    df['ds'] = pd.to_datetime(df['ds'])
    
    # 获取上个月月初日期
    last_month_input = datetime.now().replace(day=1, hour=0, minute=0, second=0, microsecond=0) - pd.DateOffset(months=1)
    
    # 定义分组检查函数
    def check_consecutive(group):
        # 提取月份维度并去重、降序排序
        sorted_month_periods = group['ds'].dt.to_period('M').drop_duplicates().sort_values(ascending=False)
        # 转换为每月第一天的datetime格式
        sorted_dates = sorted_month_periods.to_timestamp()
        
        # 如果该组最新日期早于上个月，直接不满足
        if sorted_dates.iloc[0] < last_month_input:
            return False
        
        # 计算相邻月份的间隔（取整为月）
        month_diff = (sorted_dates.shift(-1) - sorted_dates) / pd.Timedelta(days=30.44)
        month_diff = month_diff.round().fillna(0)
        
        # 统计从上个月开始的连续月份数
        consecutive_count = 0
        for date, diff in zip(sorted_dates, month_diff):
            # 跳过本月，从上个月开始检查
            if date > last_month_input:
                continue
            # 相邻月份差为-1表示连续
            if diff == -1:
                consecutive_count += 1
                if consecutive_count >= min_consecutive_observations:
                    return True
            else:
                # 中断连续则重置计数
                consecutive_count = 0
                # 若当前日期已超出需要检查的范围，提前终止
                if date < last_month_input - pd.DateOffset(months=min_consecutive_observations):
                    break
        # 最后验证是否满足最小连续数要求
        return consecutive_count >= min_consecutive_observations
    
    # 筛选符合条件的unique_id
    valid_ids = df.groupby('unique_id').apply(check_consecutive)
    valid_ids = valid_ids[valid_ids].index.tolist()
    
    # 返回筛选后的结果
    return df[df['unique_id'].isin(valid_ids)]

测试验证 #

调用函数测试：

result = filter_consecutive_ids(input_df, min_consecutive_observations=2)
print(result)

输出结果与预期一致：

ds    y  unique_id
0 2024-02-01  500          1
1 2024-01-01  600          1
2 2023-12-01  700          1

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