Collapse multiple timestamp rows into a single one

user_9112151 0
Viewed 88

I have a series like that:

s = pd.DataFrame({'ts': [1, 2, 3, 6, 7, 11, 12, 13]})
s

    ts
0   1
1   2
2   3
3   6
4   7
5   11
6   12
7   13

I would like to collapse rows that have difference less than MAX_DIFF (2). That means that the desired output must be:

[{'ts_from': 1, 'ts_to': 3},
 {'ts_from': 6, 'ts_to': 7},
 {'ts_from': 11, 'ts_to': 13}]

I did some coding:

s['close'] = s.diff().shift(-1)
s['close'] = s[s['close'] > MAX_DIFF].astype('bool')
s['close'].iloc[-1] = True

parts = []
ts_from = None

for _, row in s.iterrows():
    if row['close'] is True:
        part = {'ts_from': ts_from, 'ts_to': row['ts']}
        parts.append(part)
        ts_from = None
        continue
    
    if not ts_from:
        ts_from = row['ts']

This works but does not seem optimal because of iterrows(). I thought about ranks but couldn't figure out how to implement them so as to groupby rank further.

Is there way to optimes algorithm?

pandas python
1 Answers

You can create groups by checking where the difference is more than your threshold and take a cumsum. Then agg however you'd like, perhaps first and last in this case.

gp = s['ts'].diff().abs().ge(2).cumsum().rename(None)
res = s.groupby(gp).agg(ts_from=('ts', 'first'),
                        ts_to=('ts', 'last'))
#   ts_from  ts_to
#0        1      3
#1        6      7
#2       11     13

And if you want the list of dicts then:

res.to_dict('records')
#[{'ts_from': 1, 'ts_to': 3},
# {'ts_from': 6, 'ts_to': 7},
# {'ts_from': 11, 'ts_to': 13}]

For completeness here is how the grouper aligns with the DataFrame:

s['gp'] = gp
print(s)

   ts  gp
0   1   0     # `1` becomes ts_from for group 0
1   2   0
2   3   0     # `3` becomes ts_to for group 0
3   6   1     # `6` becomes ts_from for group 1
4   7   1     # `7` becomes ts_to for group 1
5  11   2     # `11` becomes ts_from for group 2
6  12   2
7  13   2     # `13` becomes ts_to for group 2
user_43333590
Related