фильтровать фрейм данных на основе идентификатора и диапазона дат

Question

У меня есть pandas dataframe, df, содержащий ID и date столбцы:

start = datetime.datetime.today()
dates = [start, start+relativedelta(days=20), start+relativedelta(days=40),
         start, start+relativedelta(days=35), start+relativedelta(days=36),
         start, start+relativedelta(days=10), start+relativedelta(days=15)]

df = pd.DataFrame({'ID':[1,1,1,2,2,2,3,3,3], 'date':dates})

   ID                       date
0   1 2018-11-29 15:35:56.876549
1   1 2018-12-19 15:35:56.876549
2   1 2019-01-08 15:35:56.876549
3   2 2018-11-29 15:35:56.876549
4   2 2019-01-03 15:35:56.876549
5   2 2019-01-04 15:35:56.876549
6   3 2018-11-29 15:35:56.876549
7   3 2018-12-09 15:35:56.876549
8   3 2018-12-14 15:35:56.876549

Теперь я хочу отфильтровать df, чтобы для каждого идентификатора только первые 30дни включены.Т.е. дата <= (date.min () + 30 дней) </p>

Это означает, например, что ID = 1, 2019-01-08, это более 30 дней после первой даты, 2018-11-29, поэтомуэто должно быть удалено.И так далее.Результирующий новый фрейм данных должен быть:

   ID                       date
0   1 2018-11-29 15:35:56.876549
1   1 2018-12-19 15:35:56.876549
3   2 2018-11-29 15:35:56.876549
6   3 2018-11-29 15:35:56.876549
7   3 2018-12-09 15:35:56.876549
8   3 2018-12-14 15:35:56.876549

Как это можно сделать программно?

Answer 1

Рекомендуется добавить вспомогательные столбцы для дат start и end , а затем запустить булево индексирование для фильтра.В частности, используйте groupby().tansform для встроенного min агрегирования:

df['start_date'] = df.groupby(df['ID'])['date'].transform('min')
df['end_date'] = df['start_date'] + relativedelta(days=30)

# BOOLEAN MASK
sub_df = df[(df['date'] >= df['start_date']) & (df['date'] <= df['end_date'])]
print(sub_df)
#    ID                       date                 start_date                   end_date
# 0   1 2018-11-29 15:22:35.301788 2018-11-29 15:22:35.301788 2018-12-29 15:22:35.301788
# 1   1 2018-12-19 15:22:35.301788 2018-11-29 15:22:35.301788 2018-12-29 15:22:35.301788
# 3   2 2018-11-29 15:22:35.301788 2018-11-29 15:22:35.301788 2018-12-29 15:22:35.301788
# 6   3 2018-11-29 15:22:35.301788 2018-11-29 15:22:35.301788 2018-12-29 15:22:35.301788
# 7   3 2018-12-09 15:22:35.301788 2018-11-29 15:22:35.301788 2018-12-29 15:22:35.301788
# 8   3 2018-12-14 15:22:35.301788 2018-11-29 15:22:35.301788 2018-12-29 15:22:35.301788

# WITH BETWEEN()
sub_df = df[df['date'].between(df['start_date'], df['end_date'])]
print(sub_df)
#    ID                       date                 start_date                   end_date
# 0   1 2018-11-29 15:22:35.301788 2018-11-29 15:22:35.301788 2018-12-29 15:22:35.301788
# 1   1 2018-12-19 15:22:35.301788 2018-11-29 15:22:35.301788 2018-12-29 15:22:35.301788
# 3   2 2018-11-29 15:22:35.301788 2018-11-29 15:22:35.301788 2018-12-29 15:22:35.301788
# 6   3 2018-11-29 15:22:35.301788 2018-11-29 15:22:35.301788 2018-12-29 15:22:35.301788
# 7   3 2018-12-09 15:22:35.301788 2018-11-29 15:22:35.301788 2018-12-29 15:22:35.301788
# 8   3 2018-12-14 15:22:35.301788 2018-11-29 15:22:35.301788 2018-12-29 15:22:35.301788

# WITH QUERY()
sub_df = df.query('date >= start_date & date <= end_date')
print(sub_df)
#    ID                       date                 start_date                   end_date
# 0   1 2018-11-29 15:22:35.301788 2018-11-29 15:22:35.301788 2018-12-29 15:22:35.301788
# 1   1 2018-12-19 15:22:35.301788 2018-11-29 15:22:35.301788 2018-12-29 15:22:35.301788
# 3   2 2018-11-29 15:22:35.301788 2018-11-29 15:22:35.301788 2018-12-29 15:22:35.301788
# 6   3 2018-11-29 15:22:35.301788 2018-11-29 15:22:35.301788 2018-12-29 15:22:35.301788
# 7   3 2018-12-09 15:22:35.301788 2018-11-29 15:22:35.301788 2018-12-29 15:22:35.301788
# 8   3 2018-12-14 15:22:35.301788 2018-11-29 15:22:35.301788 2018-12-29 15:22:35.301788

Для очистки вспомогательных столбцов:

# DROP HELPER COLUMNS
sub_df = sub_df.drop(columns=['start_date', 'end_date'])
print(sub_df)
#    ID                       date
# 0   1 2018-11-29 15:22:35.301788
# 1   1 2018-12-19 15:22:35.301788
# 3   2 2018-11-29 15:22:35.301788
# 6   3 2018-11-29 15:22:35.301788
# 7   3 2018-12-09 15:22:35.301788
# 8   3 2018-12-14 15:22:35.301788