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