Блог

Учитывая фрейм данных с наблюдениями, как можно возвращать строки, которые находятся в пределах -X дней от заданного списка дат?

Я придумал следующую функцию, но есть ли более простой и эффективный способ решения этой задачи?

 import numpy as np
from numpy.random import RandomState


def filterDfByDates(df, dates_of_observations, date_range):

    """
        Extract all rows in the dataframe which fall between any date in the dates_of_observation  - date_range range
    """

    ##Build mask
    mask = np.full(df.shape[0],False)

    for query_date in dates_of_observations:
        min_day = query_date - date_range
        max_day = query_date   date_range
        mask |= ( (df.index >= min_day) amp; (df.index <= max_day) )

    return df[mask]


rand = RandomState(17)
dates : np.ndarray = rand.choice(a=np.arange(np.datetime64('2021-01-01'),
                  np.datetime64('2021-01-15'), np.timedelta64(1, 'h')),size= 30, replace=True)
dates.sort()
randData = rand.choice([True, False], len(dates), p=[0.1, 0.9])
df = pd.DataFrame({"event": randData},
                  index=dates)

dates_of_obs = df.query("event").index

filterDfByDates(df,dates_of_obs, np.timedelta64(1,'D'))
 

От вашего DataFrame :

 >>> import pandas as pd
>>> from io import StringIO

>>> df = pd.read_csv(StringIO("""
date,event
2012-01-01 12:30:00,event1
2012-01-01 12:30:12,event2
2012-01-01 12:30:12,event3
2012-01-02 12:28:29,event4
2012-02-01 12:30:29,event4
2012-02-01 12:30:38,event5
2012-03-01 12:31:05,event6
2012-03-01 12:31:38,event7
2012-06-01 12:31:44,event8
2012-07-01 10:31:48,event9
2012-07-01 11:32:23,event10"""))
>>> df['date'] = pd.to_datetime(df['date'], format="%Y-%m-%d %H:%M:%S.%f")
>>> df
    date                event
0   2012-01-01 12:30:00 event1
1   2012-01-01 12:30:12 event2
2   2012-01-01 12:30:12 event3
3   2012-01-02 12:28:29 event4
4   2012-02-01 12:30:29 event4
5   2012-02-01 12:30:38 event5
6   2012-03-01 12:31:05 event6
7   2012-03-01 12:31:38 event7
8   2012-06-01 12:31:44 event8
9   2012-07-01 10:31:48 event9
10  2012-07-01 11:32:23 event10
 

Во — первых, мы начинаем со смещения date столбца и вычитаем его из исходного date столбца :

 >>> g = df['date'].sub(df['date'].shift(1)).dt.days
>>> g
0      NaN
1      0.0
2      0.0
3      0.0
4     30.0
5      0.0
6     29.0
7      0.0
8     92.0
9     29.0
10     0.0
Name: date, dtype: float64
 

Затем мы применяем a cumsum для всех значений, превышающих X (здесь это 1 день), чтобы получить ожидаемый результат :

 >>> X = 1
>>> df.groupby(g.gt(X).cumsum()).apply(print)
                 date   event
0 2012-01-01 12:30:00  event1
1 2012-01-01 12:30:12  event2
2 2012-01-01 12:30:12  event3
3 2012-01-02 12:28:29  event4
                 date   event
4 2012-02-01 12:30:29  event4
5 2012-02-01 12:30:38  event5
                 date   event
6 2012-03-01 12:31:05  event6
7 2012-03-01 12:31:38  event7
                 date   event
8 2012-06-01 12:31:44  event8
                 date   event
9 2012-07-01 10:31:48  event9
10 2012-07-01 11:32:23 event10