python - 如何将代码构造编写为函数

Question

我是编程和 python 的新手，想使用 'def' 'return' 构造将以下代码编写为函数：

df.loc[df['DATE_INT'].shift(-1) - df['DATE_INT'] == 1, 'CONSECUTIVE_DAY'] = True
df.loc[(df['DATE_INT'].shift(-1) - df['DATE_INT'] == 1) | (df['DATE_INT'].shift(1) - df['DATE_INT'] == -1), 'CONSECUTIVE_DAY'] = True

我的尝试返回无效语法：

def ConsecutiveID(df, column ='DATE_INT'):
    return  df.loc[df['DATE_INT'].shift(-1) - df['DATE_INT'] == 1, 'CONSECUTIVE_DAY'] = True
            df.loc[(df['DATE_INT'].shift(-1) - df['DATE_INT'] == 1) | (df['DATE_INT'].shift(1) - df['DATE_INT'] == -1), 'CONSECUTIVE_DAY'] = True

我的目标是最终使用我的 ConsecutiveID 函数，如下所示：

    df.groupby(['COUNTY_GEOID_YEAR','TEMPBIN']).apply(ConsecutiveID)

我正在应用拆分应用组合结构。groupby 在哪里拆分我的数据，我使用我想在 apply 中构造的函数。

我的主要问题是如何将我称之为 ConsecutiveID 的内容编写为函数。感谢您的任何帮助。

Answer 1

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

    df.loc[cond1 | cond2, 'CONSECUTIVE_DAY'] = True

    return df