基于groupby筛选最长正负连续序列并生成目标DataFrame

原始数据 #

import pandas as pd

df = pd.DataFrame(
    {
        'a': ['a', 'a', 'a', 'c', 'b', 'a', 'a', 'b', 'c'],
        'b': [20, 20, 20,-70, 70, -10, -10, -1, -1],
    }
)

目标结果 #

需要生成包含两行的DataFrame，分别提取最长的正值连续序列（标记为long）和最长的负值连续序列（标记为short）；若存在多个同长度的最长序列，选择sum值更大的那个：

direction   length  sum
0   long        3        60
1   short       2       -20

当前已完成的步骤 #

已通过以下代码给每个连续序列打上分组标记，并计算了每个序列的长度：

df['streak'] = df['b'].ne(df['b'].shift()).cumsum()
df['size'] = df.groupby('streak')['b'].transform('size')

得到的中间结果：

a   b  streak  size
0  a  20       1     3
1  a  20       1     3
2  a  20       1     3
3  c -70       2     1
4  b  70       3     1
5  a -10       4     2
6  a -10       4     2
7  b  -1       5     2
8  c  -1       5     2

后续解决方案 #

按以下步骤即可完成需求：

• 给每个序列标记方向（long/short）：根据b列值的正负判断（假设b列无0值，若有0需额外处理）
• 按streak分组，计算每个序列的sum和length
• 分别对long和short组筛选最长序列，若长度相同选sum更大的
• 整理成目标DataFrame

完整代码如下：

import pandas as pd

df = pd.DataFrame(
    {
        'a': ['a', 'a', 'a', 'c', 'b', 'a', 'a', 'b', 'c'],
        'b': [20, 20, 20,-70, 70, -10, -10, -1, -1],
    }
)

# 标记连续序列分组
df['streak'] = df['b'].ne(df['b'].shift()).cumsum()

# 计算每个序列的方向、长度、求和值
seq_stats = df.groupby('streak').agg(
    direction=('b', lambda x: 'long' if x.iloc[0] > 0 else 'short'),
    length=('b', 'size'),
    sum=('b', 'sum')
).reset_index(drop=True)

# 筛选每个方向的最优序列：先按长度降序，再按sum降序，取每组第一行
result = seq_stats.sort_values(by=['length', 'sum'], ascending=[False, False]) \
                  .groupby('direction').first().reset_index()

# 调整列顺序以匹配目标格式
result = result[['direction', 'length', 'sum']]

print(result)

运行结果：

direction  length  sum
0      long       3   60
1     short       2  -20

代码说明 #

• groupby('streak').agg()：一次性计算每个连续序列的方向、长度和求和值，方向由序列第一个元素的正负判断
• sort_values(...)：先按长度从大到小排序，长度相同则按sum从大到小排序，确保同长度序列中sum更大的排在前列
• groupby('direction').first()：对long和short分别取排序后的第一行，即符合要求的最优序列

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