Pandas - groupby continuous datetime periods

user_5805262 0
Viewed 1404

I have a pandas dataframe that looks like this:

    KEY   START       END         VALUE
0   A     2017-01-01  2017-01-16  2.1
1   B     2017-01-01  2017-01-23  4.3
2   B     2017-01-23  2017-02-10  1.7
3   A     2017-01-28  2017-02-02  4.2
4   A     2017-02-02  2017-03-01  0.8

I would like to groupby on KEY and sum on VALUE but only on continuous periods of time. For instance in the above example I would like to get:

   KEY  START       END         VALUE 
0  A    2017-01-01  2017-01-16  2.1
1  A    2017-01-28  2017-03-01  5.0
2  B    2017-01-01  2017-02-10  6.0

There are tow groups for A since there is a gap in the time periods. I would like to avoid for loops since the dataframe has tens of millions of rows.

pandas-groupby pandas python
2 Answers

Create helper Series by compare shifted START column per group and use it for groupby:

s = df.loc[df.groupby('KEY')['START'].shift(-1) == df['END'], 'END']
s = s.combine_first(df['START'])
print (s)
0   2017-01-01
1   2017-01-23
2   2017-01-23
3   2017-02-02
4   2017-02-02
Name: END, dtype: datetime64[ns]

df = df.groupby(['KEY', s], as_index=False).agg({'START':'first','END':'last','VALUE':'sum'})
print (df)
  KEY  VALUE      START        END
0   A    2.1 2017-01-01 2017-01-16
1   A    5.0 2017-01-28 2017-03-01
2   B    6.0 2017-01-01 2017-02-10
user_29010020

The answer from jezrael works like a charm if there are only two consecutive rows to aggregate. In the new example, it would not aggregate the last three rows for KEY = A.

     KEY   START       END         VALUE
 0   A     2017-01-01  2017-01-16  2.1
 1   B     2017-01-01  2017-01-23  4.3
 2   B     2017-01-23  2017-02-10  1.7
 3   A     2017-01-28  2017-02-02  4.2
 4   A     2017-02-02  2017-03-01  0.8 
 5   A     2017-03-01  2017-03-23  1.0

The following solution (slight modification of jezrael's solution) enables to aggregate all rows that should be aggregated:

df = df.sort_values(by='START')
idx = df.groupby('KEY')['START'].shift(-1) != df['END']
df['DATE'] = df.loc[idx, 'START']
df['DATE'] = df.groupby('KEY').DATE.fillna(method='backfill')
df = (df.groupby(['KEY', 'DATE'], as_index=False)
        .agg({'START': 'first', 'END': 'last', 'VALUE': 'sum'})
        .drop(['DATE'], axis=1))

Which gives:

   KEY   START         END      VALUE
0   A  2017-01-01  2017-01-16    2.1
1   A  2017-01-28  2017-03-23    6.0
2   B  2017-01-01  2017-02-10    6.0

Thanks @jezrael for the elegant approach!

user_58052620
Related