如何检查df中unique_id的记录数是否符合特定数量？

Question

我有一个这样的数据框：

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)

我需要一个函数，来选择那些在上个月往回看时，符合最小连续观察次数的唯一ID。

如果最小连续观察次数设为2，预期的输出是：

expected_output = {
    'ds': ['2024-02-01', '2024-01-01', '2023-12-01'],
    'y': [500, 600, 700],
    'unque_id': [1, 1, 1]
}

我在编写这个函数时遇到了困难。

我这样得到了上个月的数据：

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

但是我不知道怎么检查从上个月的数据往回看是否连续。

Answer 1

如果我理解正确的话，要把数据转换成按月的周期，可以使用diff这个方法来判断两个月份是否是连续的（这里的参考是MonthEnd），然后计算连续的月份有多少。如果连续的月份数量大于或等于你设定的阈值，就保留这个ID：

# min consecutive observations
N = 2

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

def min_cons_obs(g, N):
    return (g.groupby(g[::-1].diff().ne(pd.offsets.MonthEnd()).cumsum())
            .size().ge(N).any()
            )

keep = (df['ds'].dt.to_period('M')
        .groupby(df['unique_id'])
        .agg(min_cons_obs, N=N)
        )

out = df[df['unique_id'].isin(keep.index[keep])]

输出结果：

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

另外，你也可以计算每一段连续的月份数量，这样可以让你在筛选时有更多的灵活性，使用自定义的逻辑：

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

def min_cons_obs(g):
    return (g.groupby(g[::-1].diff().ne(pd.offsets.MonthEnd()).cumsum())
            .transform('size')
            )

df['consecutive'] = (df['ds']
        .dt.to_period('M')
        .groupby(df['unique_id'])
        .transform(min_cons_obs)
        )

输出结果：

          ds    y  unique_id  consecutive
0 2024-02-01  500          1            3
1 2024-01-01  600          1            3
2 2023-12-01  700          1            3
3 2024-02-01  800          2            1
4 2023-12-01  500          2            1