根据 pandas 中的顺序数据集中的标签创建间隔块

Question

所以我尝试使用我已有的每月数据集为生存分析构建数据集。 我当前使用 Pandas 的 Python 代码是这样的（我知道这段代码还不完美，但我在这里要求使用 Big O 符号）：

chrn_months = df.loc[df.label==1].month
temp_month = 0
df_final = pd.DataFrame()
for month in chrn_months:
    df_temp = df.loc[df.month==month]
    df_temp['start'] = temp_month
    df_temp['end'] = month
    df_final = pd.concat([df_final, df_temp])
    temp_month = month
    df = df.loc[df.month>=month]
if list(chrn_months)[-1] != 16:
    df_temp['start'] = df_final.end.max()
    df_temp['end'] = 16
    df_final = pd.concat([df_final, df_temp])

我已经修复了上面代码中的几乎所有故障，现在它可以工作了，我设法编写了代码，但问题是它的计算量太大，我想知道是否有更好的方法来做到这一点。 另外，我还需要总结月份间隔内其他列的值，因此请记住这一点。 正如你所看到的，我使用 for 循环为每个 ID 创建所需的输出，我们还需要使用另一个循环（嵌套）来过滤掉巨型数据框中每个 ID 的数据，这样我们就得到了 Big O( n^2) 这对我来说是有问题的。 我运行代码，几个小时后它不起作用并由于包含大量数据而冻结。 任何建议将不胜感激。

Answer 1

使用.query()、.assign()、.shift()：

数据：

import pandas as pd


data = {
    "id": [123456, 123456, 123456, 123456, 123456, 123456, 123456, 123456,
           123456, 123456, 123456, 123456, 123456, 123456, 123456, 123456],
    "month": [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16],
    "label": [1, 1, 0, 1, 1, 0, 0, 0, 1, 0, 0, 1, 0, 0, 0, 1]
}
df = pd.DataFrame(data=data)

代码：

df = df.query(expr="label.eq(1)")
df = df.assign(
    start=df.month.shift(1).fillna(0).astype(int),
    end=df.month.shift(0)
).filter(["id", "start", "end"]).reset_index(drop=True)
print(df)

输出：

       id  start  end
0  123456      0    1
1  123456      1    2
2  123456      2    4
3  123456      4    5
4  123456      5    9
5  123456      9   12
6  123456     12   16