大数据集下DataFrame性能优化：根包过滤算法提升问询

你的核心需求是从大数据集中按name分组筛选出根包（即没有父包的包，同时排除所有子包），原代码使用双重iterrows遍历，本质还是O(n²)的时间复杂度，在数据量大时效率很低。下面是两种优化方案，从简单高效到极致优化：

优化思路 #

1. 按name分组处理：仅在同组内筛选根包，避免跨组无效比较
2. 利用排序特性+根包集合：每组内按包长度升序排序后，维护已确认的根包集合，遍历每个包时只需检查是否属于集合中某个根包的子包，无需遍历所有后续元素
3. 替换iterrows：改用pandas分组操作和原生列表遍历，规避iterrows的低效问题

基础优化代码 #

import pandas as pd
import tabulate

def dumpdf(df):
    if len(df) == 0:
        return
    df = df.reset_index(drop=True)
    tab = tabulate.tabulate(df, headers='keys', tablefmt='psql', showindex=True)
    print(tab)

def filter_root_packages(group):
    # 按包长度升序排序，短包优先处理
    sorted_pkgs = group['package'].sort_values(key=lambda x: x.str.len()).tolist()
    root_pkgs = []
    for pkg in sorted_pkgs:
        # 检查当前包是否是已保留根包的子包
        is_child = any(pkg.startswith(root) for root in root_pkgs)
        if not is_child:
            root_pkgs.append(pkg)
    # 返回仅包含根包的分组数据
    return group[group['package'].isin(root_pkgs)]

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'])
    
    # 按name分组筛选根包
    root_df = df.groupby('name', group_keys=False).apply(filter_root_packages)
    
    # 按需求聚合结果
    result_df = root_df.groupby('name', as_index=False).agg({'package':'\n'.join})
    dumpdf(result_df)

if __name__ == "__main__":
    main()

极致优化：前缀树（字典树）方案 #

如果数据集极大且根包数量较多，可使用前缀树将前缀判断的时间复杂度从O(k)（k为根包数量）降至O(L)（L为包的分段数）：

import pandas as pd
import tabulate

def dumpdf(df):
    if len(df) == 0:
        return
    df = df.reset_index(drop=True)
    tab = tabulate.tabulate(df, headers='keys', tablefmt='psql', showindex=True)
    print(tab)

class TrieNode:
    def __init__(self):
        self.children = {}
        self.is_end = False

def is_child_pkg(trie_root, pkg):
    node = trie_root
    parts = pkg.split('.')
    for i, part in enumerate(parts):
        if part not in node.children:
            return False
        node = node.children[part]
        # 中途遇到已标记为根包的节点，说明当前包是子包
        if node.is_end and i < len(parts)-1:
            return True
    return False

def filter_root_packages_with_trie(group):
    sorted_pkgs = group['package'].sort_values(key=lambda x: x.str.len()).tolist()
    trie_root = TrieNode()
    root_pkgs = []
    for pkg in sorted_pkgs:
        if not is_child_pkg(trie_root, pkg):
            root_pkgs.append(pkg)
            # 将当前根包插入前缀树
            node = trie_root
            for part in pkg.split('.'):
                if part not in node.children:
                    node.children[part] = TrieNode()
                node = node.children[part]
            node.is_end = True
    return group[group['package'].isin(root_pkgs)]

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'])
    
    # 使用前缀树方案筛选根包
    root_df = df.groupby('name', group_keys=False).apply(filter_root_packages_with_trie)
    
    result_df = root_df.groupby('name', as_index=False).agg({'package':'\n'.join})
    dumpdf(result_df)

if __name__ == "__main__":
    main()

优化效果说明 #

• 原代码时间复杂度为O(n²)，基础优化版为O(nk)（k为平均根包数量），前缀树版为O(nm)（m为包的平均分段数），大数据集下性能提升显著
• 彻底避免了iterrows的低效遍历，改用pandas分组操作和原生逻辑，大幅降低内存开销与循环耗时

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