python - 识别连续的数据序列并计算其长度

Question

我正在使用一个 DataFrame，其中每行观察都附加了一个序数日期时间对象。我编写了一个函数，我相信该函数会查看我的 DataFrame 并使用以下代码识别连续发生的日期以及这些连续发生的日期的运行长度:

 def consecutiveCount(df):
    df= df.copy()
    cond1 = df['DATE_INT'].shift(-1) - df['DATE_INT'] == 1
    cond2 = df['DATE_INT'].shift(1) - df['DATE_INT'] == -1

    cond3 = df['DATE_INT'].shift(-2) - df['DATE_INT'] == 2            
    cond4 = df['DATE_INT'].shift(2) - df['DATE_INT'] == -2

现在我继续以同样的方式提出这些条件，直到这一点:

    cond55 = df['DATE_INT'].shift(-28) - df['DATE_INT'] == 28            
    cond56 = df['DATE_INT'].shift(28) - df['DATE_INT'] == -28

    cond57 = df['DATE_INT'].shift(-29) - df['DATE_INT'] == 29            
    cond58 = df['DATE_INT'].shift(29) - df['DATE_INT'] == -29

然后，我使用以下代码将“运行”天数的长度写入列变量中:

    df.loc[cond1 | cond2, 'CONSECUTIVE_COUNT'] = 2
    df.loc[cond3 | cond4, 'CONSECUTIVE_COUNT'] = 3

我再次继续，直到达到长度为 30 天的“运行”。

    df.loc[cond55 | cond56, 'CONSECUTIVE_COUNT'] = 29
    df.loc[cond57 | cond58, 'CONSECUTIVE_COUNT'] = 30

最后，我将该函数应用于 DataFrame 的特定组，如下所示:

    df1 = df.groupby(['COUNTY_GEOID_YEAR','TEMPBIN']).apply(consecutiveCount)

我确信有更有效的方法来编写这段代码。我已经确定脚本中的瓶颈是通过在整个脚本中打印各种字符串来应用该函数。

任何有关以更有效的方式编写函数或如何加快应用函数的帮助都会很棒!如果我可以提供更多信息，请告诉我。

编辑:正如@DSM指出的那样，我的代码没有准确地正确计算连续运行天数的长度。他的解决方案对我来说非常准确!

Answer 1

IIUC，您可以在应用 groupby 后使用 shift-compare-cumsum 模式，然后进行转换。

假设您的数据看起来像这样(稍微简化一下)

df = pd.DataFrame({"GEOID_YEAR": [2000]*10 + [2001]*4, "TEMPBIN": [1]*14,
                   "DATE_INT": [1,2,3,4,6,7,9,10,11,14] + list(range(14,18)),
                   "OTHER_COL": [2]*14})

或

    DATE_INT  GEOID_YEAR  OTHER_COL  TEMPBIN
0          1        2000          2        1
1          2        2000          2        1
2          3        2000          2        1
3          4        2000          2        1
4          6        2000          2        1
5          7        2000          2        1
6          9        2000          2        1
7         10        2000          2        1
8         11        2000          2        1
9         14        2000          2        1
10        14        2001          2        1
11        15        2001          2        1
12        16        2001          2        1
13        17        2001          2        1

然后

df["cons_id"] = df.groupby(["GEOID_YEAR", "TEMPBIN"])["DATE_INT"].apply(
    lambda x: (x != x.shift() + 1).cumsum())
df["cons_count"] = (df.groupby(["GEOID_YEAR", "TEMPBIN", "cons_id"])
                    ["cons_id"].transform("count"))

给我们

In [78]: df
Out[78]: 
    DATE_INT  GEOID_YEAR  OTHER_COL  TEMPBIN  cons_id  cons_count
0          1        2000          2        1        1           4
1          2        2000          2        1        1           4
2          3        2000          2        1        1           4
3          4        2000          2        1        1           4
4          6        2000          2        1        2           2
5          7        2000          2        1        2           2
6          9        2000          2        1        3           3
7         10        2000          2        1        3           3
8         11        2000          2        1        3           3
9         14        2000          2        1        4           1
10        14        2001          2        1        1           4
11        15        2001          2        1        1           4
12        16        2001          2        1        1           4
13        17        2001          2        1        1           4