Сохранять последние N записей каждой группы в pandas DataFrame

Question

У меня есть pandas df, который выглядит следующим образом.

df = pd.DataFrame({'a':['ABC', 'ABC', 'DEF', 'DEF', 'DEF', 'DEF'], 'b':['1', '1', '2', '2', '2', '2'], 'c':[0.1, 0.2, 0.3, 0.4, 0.5, 0.6]})

    a    b    c
0  ABC   1    0.1
1  ABC   1    0.2
2  DEF   2    0.3
3  DEF   2    0.4
4  DEF   2    0.5
5  DEF   2    0.6

Я хотел бы сохранить последние n строк каждой «группы» (AB C, DEF) в столбец «a» с «n» в зависимости от значения в «b». мой результат должен быть новый df, похожий на это

    a    b    c
1  ABC   1    0.2
4  DEF   2    0.5
5  DEF   2    0.6

Answer 1

Если b не всегда увеличивается cumcount

df=df[df.iloc[::-1,].groupby('a').c.cumcount().lt(df.b.astype(int))]
df
Out[268]: 
     a  b    c
1  ABC  1  0.2
4  DEF  2  0.5
5  DEF  2  0.6

Answer 2

Используйте groupby и tail:

df.groupby('a', group_keys=False).apply(
    lambda x: x.tail(int(x['b'].iloc[0])))

     a  b    c
1  ABC  1  0.2
4  DEF  2  0.5
5  DEF  2  0.6

Это предполагает, что все значения в столбце «b» для этой группы равны, поэтому первое значение в этой группе используется для вычисления хвоста.

---

Если вам нужно сохранить равное количество строк в группе, решение упрощается до

df.groupby('a').tail(n)