Python Dataframes: фильтровать данные в соответствии с условиями группировки

Question

привет, у меня есть датафрейм, как показано ниже:

ID    date          
1     01.01.2017        
1     01.01.2017        
1     01.04.2017        
2     01.01.2017        
2     01.01.2017        
2     01.02.2017       

Я хочу отфильтровать идентификаторы, для которых минимальная и максимальная разница дат составляет 3 дня.Конечный фрейм данных должен быть таким, поскольку только идентификатор 1 соответствует условию:

ID    date          
1     01.01.2017        
1     01.01.2017        
1     01.04.2017 

Спасибо.

Answer 1

Вы можете использовать Groupby.filter с пользовательской функцией lambda, чтобы проверить, составляет ли разница между максимальной датой и минимумом 3 дней:

d = datetime.timedelta(days=3)
df.groupby('ID').date.filter(lambda x: (x.max() - x.min()) == d)

ID
1   2017-01-01
1   2017-01-01
1   2017-01-04
Name: date, dtype: datetime64[ns]

Answer 2

Вы можете создать маску и затем использовать ее в качестве фильтра:

import pandas as pd

# create sample data-frame
data = [[1, '01.01.2017'], [1, '01.01.2017'], [1, '01.04.2017'],
        [2, '01.01.2017'], [2, '01.01.2017'], [2, '01.02.2017']]
df = pd.DataFrame(data=data, columns=['id', 'date'])
df['date'] = pd.to_datetime(df.date)

# create mask
mask = df.groupby('id')['date'].transform(lambda x: (x.max() - x.min()).days == 3)

# filter
result = df[mask]

print(result)

Вывод

   id       date
0   1 2017-01-01
1   1 2017-01-01
2   1 2017-01-04