Как увеличить этап по условию в питоне панда с помощью groupby

Question

У меня есть данные, поступающие из двух групп A и B. Задача состоит в том, чтобы следить за изменениями, и если изменение (скачок) больше 4, этап устанавливается выше на 1. Данные упорядочены (временные ряды).

import pandas as pd

df = pd.DataFrame({'Group': ['A', 'A', 'A', 'A', 'A', 'A', 'B', 'B', 'B', 'B', 'B', 'B'],
                    'Leap': [1, 5, 1, 1, 5, 1, 1, 3, 5, 5, 1, 1]})

# First set Stage to 1 for all:
df['Stage'] = 1

# Function to find first leap -> set Stage to two.
def setStage2(df):
    df.loc[df['Leap'] > 4, 'Stage'] = 2
    return df

# Apply function by group:
df.groupby('Group').apply(setStage2)

Это как далеко я мог бы получить.Этап должен быть инкрементным: один раз на этапе 2, не возвращаясь к 1. Вот как должен выглядеть результат:

Так какзаполнить этап?

Answer 1

Вот одно решение, сочетающее groupby и transform.

import pandas as pd

df = pd.DataFrame({
    'Group': ['A', 'A', 'A', 'A', 'A', 'A', 'B', 'B', 'B', 'B', 'B', 'B'],
    'Leap': [1, 5, 1, 1, 5, 1, 1, 3, 5, 5, 1, 1]
})

df["Stage"] = df.groupby("Group").Leap.transform(lambda x: (x > 4).cumsum()) + 1

В этом случае вы также можете использовать apply вместо transform.