pandas Как указать значения по условию для каждой группы датафрейма

Question

Имеется фрейм данных:

data = {'фрукт': ['груша','огурец','вишня', 'абрикос', 'груша','арбуз','груша','банан', 'груша', 'вишня','яблоко', 'груша', 'вишня', 'абрикос', 'груша', 'банан'],
'страна': ['россия','сша', 'россия','россия', 'франция','ЮАР','франция', 'россия', np.nan,'россия','андорра', 'франция', 'португалия', 'россия', np.nan, 'россия'],
'id': ['01','01','01','02','02','03','03','011', '011', '011','011', '6', '6', '5', '5', '5'],
'месяц': ['март','январь','январь','март', 'сентябрь','март','октябрь', 'март', 'ноябрь', 'январь','январь', 'март', 'январь', 'март', 'июль', 'апрель']        
}
dates = pd.DataFrame(data, columns = ['фрукт','страна', 'id', 'месяц'])

Нужно получить новый фрейм, в котором для каждой группы id в колонке месяц указано значение месяца, которое следует после значения "март". Если внутри группы нет других значений, но есть только "март" - оставить значения "март" для этой группы.

Как это сделать

Какой результат нужно получить:

Answer 1

def f(col, prev_val="март"):
    res = col[col.shift().eq(prev_val) & col.ne(prev_val)]
    if res.empty:
        return prev_val
    return res.iloc[0]

df["month_new"] = df["id"].map(df.groupby("id")["месяц"].apply(f))

результат:

In [60]: df
Out[60]:
      фрукт      страна   id     месяц month_new
0     груша      россия   01      март    январь
1    огурец         сша   01    январь    январь
2     вишня      россия   01    январь    январь
3   абрикос      россия   02      март  сентябрь
4     груша     франция   02  сентябрь  сентябрь
5     арбуз         ЮАР   03      март      март
6     груша     франция   03      март      март
7     банан      россия  011      март    ноябрь
8     груша         NaN  011    ноябрь    ноябрь
9     вишня      россия  011    январь    ноябрь
10   яблоко     андорра  011    январь    ноябрь
11    груша     франция    6      март    январь
12    вишня  португалия    6    январь    январь
13  абрикос      россия    5      март      июль
14    груша         NaN    5      июль      июль
15    банан      россия    5    апрель      июль

Answer 2

Ребята, вы что, на одной овощебазе все работаете? Был уже такой вопрос, https://ru.stackoverflow.com/questions/1242489/срез-строк-внутри-группы-с-условием

БЛОГ НА HUSL

pandas Как указать значения по условию для каждой группы датафрейма

Ответы (2 шт):