Панды: удаление строк в группе на основе условия

Question

Давайте предположим, что у меня есть следующий фрейм данных:

df = pd.DataFrame({"id": [1, 1, 1, 2, 2, 2], "date": [pd.Timestamp(2016, 7, 29), pd.Timestamp(2017, 8, 22), pd.Timestamp(2017, 10, 9), pd.Timestamp(2018, 1, 9), pd.Timestamp(2018, 3, 31), pd.Timestamp(2018, 7, 5)], "other_col": [11.1, 77.7, 22.2, 33.3, 44.4, 88.8]})

Чего я хочу добиться, это удалить последние N строки для каждой группы, где data-frame is grouped by id такое, что N: Количество строк, которые имеют date, что в течение 3 месяцев с date последней строки (Это удалит последнюю строку).Результирующий фрейм данных должен быть:

r_df = pd.DataFrame({"id": [1, 2, 2], "date": [pd.Timestamp(2016, 7, 29), pd.Timestamp(2018, 1, 9), pd.Timestamp(2018, 3, 31)], "other_col": [11.1, 33.3, 44.4]})

• Строки фрейма данных всегда будут упорядочены
• Последняя строка также должна быть удалена, поскольку по умолчанию ее дата равнав пределах диапазона «удалить» даты для каждой группы.

Answer 1

Эту проблему может решить одна лямбда-функция на основе даты.Он даже не заботится о сортировке строк.

df.groupby('id').apply(
  lambda x: x.loc[
    x.date < (x.date.max() - pd.Timedelta(3, unit='M'))
  ]
).reset_index(drop=True)

Обновление :

Как указывал @jpp, следующий выбор выполняется намного (~ 12x) быстрее:

df[df['date'] < (df.groupby('id')['date'].transform('max') - pd.Timedelta(3, unit='M'))]

В качестве дополнительного бонуса также сохраняются исходный порядок и индексы.

Answer 2

def remove_n_row(grouped_df):
  last_row_date = grouped_df.iloc[-1]['date']
  grouped_df_filtered = grouped_df[(grouped_df['date'] <=last_row_date) & (grouped_df['date'] <=last_row_date - timedelta(days=90))]
  return grouped_df_filtered

from datetime import timedelta
r_df = df.groupby('id').apply(remove_n_row).reset_index(drop=True)
print(r_df)

ВЫХОД

    id     date     other_col
0   1   2016-07-29   11.1
1   2   2018-01-09   33.3
2   2   2018-03-31   44.4