Блог

У меня есть такая серия:

 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
 

Я хотел бы свернуть строки, разница между которыми меньше, чем MAX_DIFF (2) . Это означает, что желаемый результат должен быть:

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

Я сделал некоторое кодирование:

 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']
 

Это работает, но не кажется оптимальным из-за iterrows() . Я думал о рангах, но не мог понять, как их реализовать, чтобы группировать по рангу дальше.

Есть ли способ оптимизировать алгоритм?

Вы можете создавать группы, проверяя, где разница превышает ваш порог, и получать итоговую сумму. Тогда действуйте так, как вам хотелось бы, возможно first , и last в этом случае.

 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
 

И если вам нужен список dicts, тогда:

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

Для полноты картины вот как группировщик выравнивается с фреймом данных:

 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