CODEWITHSUNDEEP
pythonpandasdataframetime-series
Giang Do
Giang Do

Reputation: 93

Panda - time series shift by group

I have a dataframe like this:

pd.DataFrame({'group': {0: 1, 1: 1, 2: 1, 3: 1, 4: 2, 5: 2, 6: 2}, 'year': {0: 2007, 1: 2008, 2: 2009, 3: 2010, 4: 2006, 5: 2007, 6: 2008}, 'amount': {0: 2.0, 1: 4.0, 2: np.nan, 3: 7.0, 4: 8.0, 5: 10.0, 6: 12.0}}])

   group    year    amount
0   1       2007    2
1   1       2008    4
2   1       2009    NaN
3   1       2010    7
4   2       2006    8
5   2       2007    10
6   2       2008    12

I want to add min (remove NA), max (remove NA), lag 1 and lag 2 for each group. My ideal dataframe looks like this

   group    year    amount    min    max   lag1    lag2
0   1       2007    2          2      7     NaN     NaN
1   1       2008    4          2      7     2       NaN
2   1       2009    NaN        2      7     4       2
3   1       2010    7          2      7     NaN     4
4   2       2006    8          8      12    NaN     NaN
5   2       2007    10         8      12    8       NaN
6   2       2008    12         8      12    10      8

I know to do for min and max with agg in panda but I am not sure how to handle the lags by group

Upvotes: 1

Views: 590

Answers (1)

rafaelc
rafaelc

Reputation: 59274

IIUC

g = df.groupby('group').amount
df['max'] = g.transform('max')
df['min'] = g.transform('min')
df['lag1'] = g.shift()
df['lag2'] = g.shift(2)

group   year    amount  max     min lag1    lag2
0   1   2007    2.0     7.0     2.0 NaN     NaN
1   1   2008    4.0     7.0     2.0 2.0     NaN
2   1   2009    NaN     7.0     2.0 4.0     2.0
3   1   2010    7.0     7.0     2.0 NaN     4.0
4   2   2006    8.0     12.0    8.0 NaN     NaN
5   2   2007    10.0    12.0    8.0 8.0     NaN
6   2   2008    12.0    12.0    8.0 10.0    8.0

Upvotes: 4

Related Questions