基于groupby方法移除公共子串的代码优化需求

需求说明 #

需要移除DataFrame中存在公共子串的重复项，其中公共子串为每个name分组内的最短字符串——即保留组内最短的基准字符串，移除所有以它为前缀的长字符串。

原始实现代码 #

import re
import pandas as pd
import tabulate

def dumpdf(df):
    s = tabulate.tabulate(df, tablefmt='plain', headers='keys', showindex=True)
    print(s)
    return

def reformat(df):
    dfg = df.groupby('name')
    flag = False
    for name,df in dfg:
        item = []
        df = df.sort_values(by="package", key=lambda x: x.str.len())
        data = []
        for idx,row in df.iterrows():
            pkg = row['package']
            df.loc[df['package'].str.startswith(pkg, na=False), 'package'] = pkg
        df = df[df.duplicated(['package'], keep='first') == False]
        df = df.reset_index(drop=True)
        if len(df) > 0:
            if flag == False:
                flag = True
                out = df
            else:
                out = pd.concat([out,df],ignore_index=True)
    return out

def main():

    data = [
        ['A','com.example'],
        ['A','com.example.a'],
        ['A','com.example.b.c'],
        ['A','com.fun'],
        ['B','com.demo'],
        ['B','com.demo.b.c'],
        ['B','com.fun'],
        ['B','com.fun.e'],
        ['B','com.fun.f.g']
        ]
    df = pd.DataFrame(data,columns=['name','package'])

    df = reformat(df)
    df = df.groupby('name', as_index=False).agg('\n'.join)

    dumpdf(df)
    return

main()

原始输出 #

name    package
 0  A       com.fun
            com.example
 1  B       com.fun
            com.demo

优化后的代码 #

针对原始代码中冗余变量、循环修改DataFrame效率低、结果拼接繁琐等问题，优化后的实现如下：

import pandas as pd
import tabulate

def dumpdf(df):
    s = tabulate.tabulate(df, tablefmt='plain', headers='keys', showindex=True)
    print(s)
    return

def reformat_optimized(df):
    group_results = []
    
    for name, group in df.groupby('name'):
        # 按package长度升序排序，确保最短字符串在前
        sorted_group = group.sort_values(by='package', key=lambda x: x.str.len()).reset_index(drop=True)
        keep = [True] * len(sorted_group)
        
        # 遍历排序后的组，标记需要移除的项
        for i in range(len(sorted_group)):
            if not keep[i]:
                continue
            base_pkg = sorted_group.loc[i, 'package']
            # 后续项都是更长的字符串，直接检查是否以当前基准为前缀
            for j in range(i+1, len(sorted_group)):
                if sorted_group.loc[j, 'package'].startswith(base_pkg):
                    keep[j] = False
        
        # 筛选保留的行并加入结果列表
        filtered_group = sorted_group[keep].reset_index(drop=True)
        group_results.append(filtered_group)
    
    # 一次性合并所有组结果，避免循环concat的性能损耗
    return pd.concat(group_results, ignore_index=True)

def main():
    data = [
        ['A','com.example'],
        ['A','com.example.a'],
        ['A','com.example.b.c'],
        ['A','com.fun'],
        ['B','com.demo'],
        ['B','com.demo.b.c'],
        ['B','com.fun'],
        ['B','com.fun.e'],
        ['B','com.fun.f.g']
    ]
    df = pd.DataFrame(data, columns=['name','package'])

    df = reformat_optimized(df)
    df = df.groupby('name', as_index=False).agg('\n'.join)

    dumpdf(df)
    return

main()

优化亮点 #

• 移除冗余变量：删除原始代码中未使用的item、data变量，简化代码结构
• 避免修改原数据：通过标记keep列表筛选保留项，替代修改package列后去重的间接逻辑，逻辑更直观
• 提升拼接效率：用列表收集每个组的结果，最后一次性concat，比循环拼接DataFrame更高效
• 逻辑更清晰：直接针对排序后的组进行前缀检查，步骤明确易维护

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