基于分组与条件填充Pandas DataFrame的Code列问题

原始数据 #

创建初始DataFrame的代码：

import pandas as pd

df = pd.DataFrame({'Make': ['Tesla','Tesla','Tesla','Toyota','Ford','Ford','Ford','BMW','BMW','BMW','Mercedes','Mercedes','Mercedes'],
                   'Type': ['Model X','Model X','Model X','Corolla','Bronco','Bronco','Mustang','3 Series','3 Series','7 Series','C-Class','C-Class','S-Class'],
                   'Year': [2015, 2015, 2015, 2017, 2018, 2018, 2020, 2015, 2015, 2017, 2018, 2018, 2020],
                   'Price': [85000, 90000, 95000, 20000, 35000, 35000, 45000, 40000, 40000, 65000, 50000, 50000, 75000],
                   'Color': ['White','White','White','Red','Blue','Blue','Yellow','Silver','Silver','Black','White','White','Black'],
                   'Code'  : ['TSLABG','TSLA',None,'TYTA','FRD','_BG',None,None,'BMW','BMW','MercedesBG','Mercedes_BG','MercedesBG']
                  })

初始数据展示：

Make      Type  Year  Price   Color         Code
0    Tesla   Model X  2015  85000   White       TSLABG
1    Tesla   Model X  2015  90000   White         TSLA
2    Tesla   Model X  2015  95000   White         None
3   Toyota   Corolla  2017  20000     Red         TYTA
4     Ford    Bronco  2018  35000    Blue          FRD
5     Ford    Bronco  2018  35000    Blue         _BG
6     Ford   Mustang  2020  45000  Yellow         None
7      BMW  3 Series  2015  40000  Silver         None
8      BMW  3 Series  2015  40000  Silver          BMW
9      BMW  7 Series  2017  65000   Black          BMW
10 Mercedes  C-Class  2018  50000   White  MercedesBG
11 Mercedes  C-Class  2018  50000   White  Mercedes_BG
12 Mercedes  S-Class  2020  75000   Black  MercedesBG

需求说明 #

• 当Code列为None时，用同Make分组下的其他Code值填充
• 若同Make分组内存在带BG或_BG后缀的Code值，该分组下所有Code需统一使用对应后缀（优先保留原格式，同时移除不必要的下划线）
• 自动去除Code值中BG前的冗余下划线

尝试的代码及问题 #

尝试的代码：

code = (df['Code'].str.split('(_)', expand=True).add_prefix('part').replace('-', None).groupby(df['Make']).transform('first').fillna('').agg(''.join, axis=1))
df['Code'] = code

问题：运行后Mercedes的Code列出现MercedesBG_BG，不符合预期的MercedesBG。

正确解决方案 #

核心思路 #

按Make分组处理，先提取每个分组的基础代码和后缀类型，再统一填充并格式化Code值，避免重复添加后缀。

完整代码 #

import pandas as pd

def process_make_group(group):
    # 过滤当前分组的非空Code值
    valid_codes = group['Code'].dropna()
    if valid_codes.empty:
        return group
    
    base_codes = []
    suffix_type = None  # 记录后缀类型：None / 'BG' / '_BG'
    
    for code in valid_codes:
        # 识别后缀并提取基础代码
        if code.endswith('_BG'):
            base = code[:-3]
            suffix_type = '_BG'
        elif code.endswith('BG'):
            base = code[:-2]
            # 仅当未识别过后缀时，设置为BG
            if suffix_type is None:
                suffix_type = 'BG'
        elif code == '_BG':
            base = ''
            suffix_type = '_BG'
        else:
            base = code
        
        base_codes.append(base)
    
    # 取最长的基础代码（避免空值或短代码）
    final_base = max(base_codes, key=lambda x: len(x)) if base_codes else ''
    # 确定最终后缀
    suffix = suffix_type if suffix_type is not None else ''
    
    # 处理基础代码已包含BG的情况，避免重复添加
    if final_base.endswith('BG') and suffix in ['BG', '_BG']:
        final_base = final_base[:-2]
    
    # 为分组内所有行设置统一Code
    group['Code'] = final_base + suffix
    return group

# 按Make分组处理
df = df.groupby('Make', group_keys=False).apply(process_make_group)

正确输出结果 #

Make      Type  Year  Price   Color         Code
0    Tesla   Model X  2015  85000   White       TSLABG
1    Tesla   Model X  2015  90000   White       TSLABG
2    Tesla   Model X  2015  95000   White       TSLABG
3   Toyota   Corolla  2017  20000     Red         TYTA
4     Ford    Bronco  2018  35000    Blue       FRD_BG
5     Ford    Bronco  2018  35000    Blue       FRD_BG
6     Ford   Mustang  2020  45000  Yellow       FRD_BG
7      BMW  3 Series  2015  40000  Silver          BMW
8      BMW  3 Series  2015  40000  Silver          BMW
9      BMW  7 Series  2017  65000   Black          BMW
10 Mercedes  C-Class  2018  50000   White  MercedesBG
11 Mercedes  C-Class  2018  50000   White  MercedesBG
12 Mercedes  S-Class  2020  75000   Black  MercedesBG

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