Как отобрать срез строк со смещением?

Question

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

import pandas as pd
import numpy as np
data = {'фрукт': ['авокадо','груша','арбуз', 'груша', 'вишня', 'груша', 'банан'],
'страна': ['китай','сша', np.nan, np.nan,'канада',  np.nan, 'перу'],
'id': ['01','01','01', '011', '011',  '5', '5'],
'месяц': ['март','январь','январь', 'февраль', 'февраль', 'апрель', 'апрель']
}
dates = pd.DataFrame(data, columns = ['фрукт','страна', 'id', 'месяц'])

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

Ожидаемый результат:

Answer 1

res = (df
       .assign(страна=df["страна"].shift(), месяц=df["месяц"].shift())
       .groupby("id")
       .last()
       .reset_index())

или более гибкий вариант

cols = ["страна", "месяц"]

res = (df
       .drop(columns=cols)
       .join(df[cols].shift())
       .groupby("id")
       .last()
       .reset_index())

результат:

In [92]: res
Out[92]: 
    id  фрукт  страна    месяц
0   01  арбуз     сша   январь
1  011  вишня     NaN  февраль
2    5  банан  канада   апрель

или так если нужно именно предпоследние значения, включая NaN's:

In [108]: res = (df
     ...:        .assign(страна=df["страна"].shift(), месяц=df["месяц"].shift())
     ...:        .groupby("id", dropna=False)
     ...:        .nth(-1)
     ...:        .reset_index())

In [109]: res
Out[109]: 
    id  фрукт страна    месяц
0   01  арбуз    сша   январь
1  011  вишня    NaN  февраль
2    5  банан    NaN   апрель

БЛОГ НА HUSL

Как отобрать срез строк со смещением?

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