Обработка данных на Python: как группировать по идентификатору, а затем подмножество строк вперед / назад на 1 месяц исходя из условий времени данных в каждой группе?

Question

Я хочу установить подмножество данных после группировки, и подмножество строк вперед и назад на 1 месяц по условиям времени данных.

Вот тестовый фрейм данных:

import pandas as pd
from itertools import chain


df_1 = pd.DataFrame({
        'ID' : list(chain.from_iterable([['A'] * 365, ['B'] * 365, ['C'] * 365])),
        'Date' : pd.date_range(start = '2018-01-01', end = '2018-12-31').tolist() + pd.date_range(start = '2018-01-01', end = '2018-12-31').tolist() + pd.date_range(start = '2018-01-01', end = '2018-12-31').tolist(),
        'Value' : np.random.randn(365 * 3)
        })

df_2 = pd.DataFrame({
        'ID' : ['A', 'B', 'C'],
        'Initial_date' : pd.to_datetime(['2018-04-24', '2018-09-22', '2018-12-12'], format = '%Y-%m-%d')
        })


df = pd.merge(df_1, df_2, how = 'left', left_on = 'ID', right_on = 'ID')

вывод тестового фрейма данных:

  ID       Date  Value Initial_date
0  A 2018-01-01 -1.084   2018-04-24
1  A 2018-01-02  0.585   2018-04-24
2  A 2018-01-03  0.812   2018-04-24
3  A 2018-01-04 -1.115   2018-04-24
4  A 2018-01-05 -1.724   2018-04-24

Итак, вот что я хочу получить:

Для пользователя A его initial date равен 2018-04-24, поэтому я хочу, чтобы все строки начинались с 1 месяца.назад, что на 2018-03-24, на 1 месяц вперед, что на 2018-05-24.А также B и C пользователей.

Результат должен совпадать с

df_a = df[(df['ID'] == 'A') & (df['Date'] > '2018-03-23') & (df['Date'] < '2018-05-25')]

df_b = df[(df['ID'] == 'B') & (df['Date'] > '2018-08-21') & (df['Date'] < '2018-10-23')]

df_c = df[(df['ID'] == 'C') & (df['Date'] > '2018-11-11') & (df['Date'] < '2018-12-31')]


df_result = pd.concat([df_a, df_b, df_c])

df_result.head() равно

   ID       Date  Value Initial_date
82  A 2018-03-24 -1.013   2018-04-24
83  A 2018-03-25 -0.023   2018-04-24
84  A 2018-03-26 -0.053   2018-04-24
85  A 2018-03-27 -1.091   2018-04-24
86  A 2018-03-28  1.839   2018-04-24

Примечание *1: Для предельной даты она не обязательно включает или исключает точную дату.Например, для A, от 2018-03-23 или 2018-03-24 для меня не имеет значения.

Примечание * 2: Для C пользователей дата пересылки составляет менее 1 месяца с момента его initial date включен 2018-12-12, поэтому до конца 2018 года все в порядке.

Answer 1

IIUC, вы можете попробовать следующее:

# set 1-month interval
delta = pd.DateOffset(months=1) 

# get the subset
df_result = df[df.Date.gt(df.Initial_date-delta) & df.Date.lt(df.Initial_date+delta)]