Logo Questions Linux Laravel Mysql Ubuntu Git Menu
 

pandas: groupby and calculate time difference from first element in each group

In pandas, I would like to group data by the values in a column and then calculate the time difference between each timestamp and the first timestamp in that group.

For example, consider the following DataFrame:

# Create data. 
d = {'foo': ['001', '001', '002', '002', '002'], 
     'timestamp': ['2015-02-24 19:12:00', '2015-02-24 21:38:00', '2015-02-25 03:41:00', '2015-02-25 03:44:00', '2015-02-25 03:49:00']}
df = pd.DataFrame(d, columns = ['foo', 'timestamp'])
df['timestamp'] = pd.DatetimeIndex(pd.to_datetime(df['timestamp'])).tz_localize('UTC')
>>> print df
   foo                 timestamp
0  001 2015-02-24 19:12:00+00:00
1  001 2015-02-24 21:38:00+00:00
2  002 2015-02-25 03:41:00+00:00
3  002 2015-02-25 03:44:00+00:00
4  002 2015-02-25 03:49:00+00:00

The desired output would be:

   foo                 timestamp    output
0  001 2015-02-24 19:12:00+00:00       NaT
1  001 2015-02-24 21:38:00+00:00  02:26:00
2  002 2015-02-25 03:41:00+00:00       NaT
3  002 2015-02-25 03:44:00+00:00  00:03:00
4  002 2015-02-25 03:49:00+00:00  00:08:00

The use of .diff() gets the following, but not the desired result.

>>> d.groupby('foo')['timestamp'].diff()
0        NaT
1   02:26:00
2        NaT
3   00:03:00
4   00:05:00
like image 726
Adam Avatar asked Dec 19 '22 09:12

Adam


1 Answers

Use assign + apply

df.assign(output=df.groupby('foo').timestamp.apply(lambda x: x - x.iloc[0]))

   foo                 timestamp   output
0  001 2015-02-24 19:12:00+00:00 00:00:00
1  001 2015-02-24 21:38:00+00:00 02:26:00
2  002 2015-02-25 03:41:00+00:00 00:00:00
3  002 2015-02-25 03:44:00+00:00 00:03:00
4  002 2015-02-25 03:49:00+00:00 00:08:00
like image 63
piRSquared Avatar answered Dec 28 '22 07:12

piRSquared