Как написать конструкцию кода как функцию

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