如何筛选DataFrame中最长连续负数分组并按b列总和择优？

筛选DataFrame中符合以下规则的行：

• 选出a列中最长的连续负数序列对应的分组；
• 若存在多个长度相同的最长连续负数分组，选择b列总和更大的分组；
• 需处理DataFrame中无负数行的边界情况。

示例数据1 #

import pandas as pd
df = pd.DataFrame(
    {
        'a': [-3, -1, -2, -5, 10, -3, -13, -3, -2, 1, 2, -100],
        'b': [1, 2, 3, 4, 5, 10, 80, 90, 100, 99, 1, 12]
    }
)

期望输出1 #

a    b
5    -3   10
6   -13   80
7    -3   90
8    -2  100

逻辑规则 #

• a) 筛选a列中最长的连续负数序列对应的分组；
• b) 若多个最长连续负数分组长度相同，选择b列总和更大的分组（示例1中两个长度为4的分组，选第二个）。

尝试代码 #

import numpy as np
s = np.sign(df['a'])
df['g'] = s.ne(s.shift()).cumsum()
df['size'] = df.groupby('g')['g'].transform('size')
df['b_sum'] = df.groupby('g')['b'].transform('sum')

补充示例 #

为明确需求，另一个示例：

df = pd.DataFrame(
    {
        'a': [-0.65, 11, 18, 1, -2, -3],
        'b': [1, 20, 30000, 4322, 300, 3]
    }
)

期望输出2 #

4  -2.00    300
5  -3.00      3

现有代码问题 #

自行编写的代码可实现需求，但当DataFrame中无负数行时会报错：

df['sign'] = np.sign(df.a)
df['g'] = df.sign.ne(df.sign.shift()).cumsum()
df = df.loc[df.a.lt(0)]
out = df[df.g.eq(df.groupby('g')['b'].agg(['size', 'sum'])
                .query('size == size.max()')['sum'].idxmax())]

修正后的解决方案 #

import pandas as pd
import numpy as np

def select_longest_neg_group(df):
    # 复制原数据避免修改原DataFrame
    df_copy = df.copy()
    # 标记负数行
    df_copy['is_neg'] = df_copy['a'] < 0
    # 无负数行时返回空DataFrame
    if not df_copy['is_neg'].any():
        return pd.DataFrame(columns=df.columns)
    
    # 生成连续分组：仅对负数行进行连续分组，非负数行不参与后续计算
    df_copy['g'] = (df_copy['is_neg'] != df_copy['is_neg'].shift()).cumsum()
    # 筛选出负数行
    neg_df = df_copy[df_copy['is_neg']].copy()
    
    # 计算每个分组的长度和b列总和
    group_stats = neg_df.groupby('g')['b'].agg(['size', 'sum'])
    # 找到最长长度
    max_size = group_stats['size'].max()
    # 筛选出所有最长长度的分组，取sum最大的那个分组的g值
    target_g = group_stats[group_stats['size'] == max_size]['sum'].idxmax()
    
    # 返回目标分组的原数据列
    return neg_df[neg_df['g'] == target_g][['a', 'b']]

# 测试示例1
df1 = pd.DataFrame(
    {
        'a': [-3, -1, -2, -5, 10, -3, -13, -3, -2, 1, 2, -100],
        'b': [1, 2, 3, 4, 5, 10, 80, 90, 100, 99, 1, 12]
    }
)
print(select_longest_neg_group(df1))

# 测试示例2
df2 = pd.DataFrame(
    {
        'a': [-0.65, 11, 18, 1, -2, -3],
        'b': [1, 20, 30000, 4322, 300, 3]
    }
)
print(select_longest_neg_group(df2))

# 测试无负数的情况
df3 = pd.DataFrame({'a': [1,2,3], 'b': [4,5,6]})
print(select_longest_neg_group(df3))

代码说明 #

1. 先判断是否存在负数行，避免后续报错；
2. 仅针对负数行生成连续分组，排除非负数行的干扰；
3. 分组统计后，先筛选最长长度的分组，再从中选出b列总和最大的分组；
4. 返回结果仅保留原数据的a、b列，符合输出要求。

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