基于Pandas实现最长正负连续序列统计及首尾值提取

代码实现 #

import pandas as pd
import numpy as np

df = pd.DataFrame(
    {
        'a': [
            'a', 'a', 'a', 'a', 'a', 'a', 'a', 'a', 'a', 'a', 'a', 'a',
            'b', 'b', 'b', 'b', 'b', 'b', 'b', 'b', 'b', 'b', 'b',
        ],
        'b': [
            -20, 20, 20, 20,-70, -70, 10, -1000, -10, 100, 100,
            -11, -100, -1, -1, -100, 100, 1, 90, -1, -2, 1000, 900
        ],
        'c': [
            'f', 'f', 'f', 'f', 'f', 'x', 'x', 'x', 'y', 'y', 'y', 'a',
            'k', 'k', 'k', 'k', 'k', 't', 't', 't', 't', 's', 'e',
        ],
    }
)

# 生成连续序列分组标识与正负方向
df['direction'] = np.sign(df['b'])
df['streak'] = df['direction'].ne(df['direction'].shift()).cumsum()

# 一次性聚合所有所需指标
agg_df = df.groupby(['a', 'direction', 'streak'], as_index=False).agg(
    length=('b', 'count'),
    sum=('b', 'sum'),
    start=('c', 'first'),
    end=('c', 'last')
)

# 筛选每个(a, direction)组中长度最长的序列
out = agg_df.loc[
    agg_df.groupby(['a', 'direction'])['length'].idxmax(),
    ['a', 'direction', 'length', 'sum', 'start', 'end']
].sort_values(['a', 'direction']).reset_index(drop=True)

print(out)

代码说明 #

1. 序列分组：通过direction（正负标识）和streak（连续序列编号），将连续同符号的b值划分为同一组，避免多次重复计算。
2. 一次性聚合：在一次groupby操作中完成长度统计、求和、起始/结束值提取，逻辑更紧凑，效率远高于先聚合再合并的方式。
3. 筛选最长序列：利用idxmax()直接定位每个(a, direction)组中长度最长的序列行，避免冗余的排序和合并操作，同时保证结果准确性。

输出结果 #

a  direction  length   sum start end
0  a         -1       2 -1010     x   y
1  a          1       3    60     f   f
2  b         -1       4  -202     k   k
3  b          1       3   191     k   t

内容的提问来源于stack exchange，提问作者AmirX