如何优化DataFrame分组中多列多行求和生成is_m_days的方案？

需求说明 #

• 分组规则：以ATEXT列中的CT/RT组为核心，每组包含1行或多行随机混合的CT/RT记录，且每组前后各有一行空字符串（''）记录
• 操作目标：为每个分组对特定行的若干列进行求和，将结果输出到新列is_m_days中

原始DataFrame数据 #

import numpy as np
import pandas as pd

data = {'ATEXT': ['', 'CT', 'RT', '', '', '', '', 'CT', 'CT', 'CT', 'TT', ''], 
        'BEGUZ_UE': [11.0, 23.0, 33.0, 15.0, 12.75, 19.75, 14.75, 23.0, 
                     24.0, 24.0, 33.0, 15.0], 
        'subtract': [0.0, 0.0, 0.0, 0.2, np.nan, np.nan, 2.0, np.nan, 
                     np.nan, np.nan, np.nan, 0.0], 
        'add': [3.92, 0.0, 0.0, 0.0, np.nan, np.nan, 0.0, np.nan, np.nan, 
                np.nan, np.nan, 3.57], 
        'UE_more_days': [np.nan, np.nan, 56.0, np.nan, np.nan, np.nan, np.nan, 
                         np.nan, np.nan, np.nan, 104.0, np.nan]}
df = pd.DataFrame(data)

预期结果 #

ATEXT   BEGUZ_UE    subtract      add      UE_more_days  is_m_days
0             11.00     *0.00*        *3.92*
1     CT      *23.00*    0.00         0.00
2     RT      *33.00*    0.00         0.00          56.0
3             *15.00*    0.20         0.00                      *74.92*
4             12.75         
5             19.75
6             14.75     *2.00*       *0.00*
7     CT      *23.00*
8     RT      *24.00*
9     CT      *24.00*
10    CT      *33.00*                              104.0
11            *15.00*    0.00         3.57                     *117.00*
12
etc

现有实现代码 #

m = df['ATEXT'].eq("")
cond = (~m) & m.shift(-1)
df['UE_more_days'] = (df['BEGUZ_UE'].mask(m)
                      .groupby(m.cumsum()).cumsum()
                      .where(cond)
                     )
tmv = (df[['subtract', 'add']]
       .shift()
       .groupby(m.cumsum())
       .transform('max')
       .eval('add-subtract')
      )


df['is_m_days'] = (df.groupby(m[::-1].cumsum())['BEGUZ_UE']
                .transform('sum')
                .add(tmv)
                .where(cond)
                .shift()
               )

优化解决方案 #

现有代码逻辑可行，但可以通过更清晰的分组标识、减少中间变量、合并操作来提升可读性和效率，以下是优化后的实现：

import numpy as np
import pandas as pd

# 初始化数据
data = {'ATEXT': ['', 'CT', 'RT', '', '', '', '', 'CT', 'CT', 'CT', 'TT', ''], 
        'BEGUZ_UE': [11.0, 23.0, 33.0, 15.0, 12.75, 19.75, 14.75, 23.0, 
                     24.0, 24.0, 33.0, 15.0], 
        'subtract': [0.0, 0.0, 0.0, 0.2, np.nan, np.nan, 2.0, np.nan, 
                     np.nan, np.nan, np.nan, 0.0], 
        'add': [3.92, 0.0, 0.0, 0.0, np.nan, np.nan, 0.0, np.nan, np.nan, 
                np.nan, np.nan, 3.57], 
        'UE_more_days': [np.nan, np.nan, 56.0, np.nan, np.nan, np.nan, np.nan, 
                         np.nan, np.nan, np.nan, 104.0, np.nan]}
df = pd.DataFrame(data)

# 标记空行
m_empty = df['ATEXT'].eq('')
# 生成分组ID：每个非空组（CT/RT/TT）及前后空行归为同一分组
group_id = m_empty[::-1].cumsum()[::-1]

# 计算UE_more_days：非空行的BEGUZ_UE累加，仅保留每组最后一行非空行的值
df['UE_more_days'] = df['BEGUZ_UE'].where(~m_empty).groupby(group_id).cumsum()
df['UE_more_days'] = df['UE_more_days'].where((~m_empty) & m_empty.shift(-1, fill_value=True))

# 计算tmv：每组前导空行的(add - subtract)值
tmv = df[['add', 'subtract']].shift().groupby(group_id).first()
tmv = tmv['add'] - tmv['subtract']
tmv = tmv.reindex(group_id).values

# 计算is_m_days：每组非空行的BEGUZ_UE总和 + tmv，结果放在每组的后导空行
sum_beguz = df['BEGUZ_UE'].where(~m_empty).groupby(group_id).sum()
sum_beguz = sum_beguz.reindex(group_id).values
df['is_m_days'] = sum_beguz + tmv
df['is_m_days'] = df['is_m_days'].where(m_empty & (~m_empty).shift())

优化点说明 #

1. 分组标识更直观：通过反向累加空行标识生成group_id，直接将每个非空组及其前后空行归为同一分组，逻辑更清晰
2. 减少重复操作：合并部分分组计算步骤，避免多次重复分组，提升运行效率
3. 可读性提升：每个步骤添加明确注释，操作目标清晰，后续维护更方便
4. 逻辑更严谨：使用first()替代max()获取前导空行的add/subtract值，更符合“每组前一行数据”的需求定义（原代码max()在多空行场景可能出现错误）

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