性能优化：统计DataFrame中以特定字符串结尾的唯一行数量

核心思路 #

原代码通过apply逐行执行全表扫描，时间复杂度为O(N²)，对100万行数据效率极低。优化核心是一次性预处理所有后缀对应的唯一字符串计数，再通过字典映射快速填充结果，时间复杂度降至O(M)（M为column_A的唯一值数量），能大幅缩短运行时间。

具体实现代码 #

import pandas as pd

# 1. 提取column_A的唯一值，避免重复计算
unique_vals = df['column_A'].unique()

# 2. 筛选含"."的唯一值（只有这类值会被某个s的统计纳入）
valid_vals = [val for val in unique_vals if '.' in val]

# 3. 生成（后缀s_candidate, 原字符串）配对列表
#    后缀取原字符串最后一个"."之后的部分，对应能匹配该后缀的s
suffix_pairs = [(val.split('.')[-1], val) for val in valid_vals]

# 4. 按后缀分组，统计每个后缀对应的唯一原字符串数量
suffix_counts = pd.DataFrame(suffix_pairs, columns=['suffix', 'val'])\
    .groupby('suffix')['val'].nunique()

# 5. 映射统计结果到原DataFrame，无匹配的填充0
df['column_B'] = df['column_A'].map(suffix_counts.to_dict()).fillna(0).astype(int)

关键优化点说明 #

• 去重预处理：先处理column_A的唯一值，避免对重复的s重复执行统计逻辑。
• 定向筛选：仅处理包含.的字符串，排除无法被任何s匹配的无效数据。
• 分组统计：通过一次分组完成所有后缀的计数，替代逐行全表扫描的低效操作。
• 快速映射：用字典映射替代apply，填充结果的时间复杂度为O(N)，远快于原方法。

边界情况处理 #

• 若某个s没有任何字符串以.{s}结尾，map返回的NaN会被fillna(0)填充为0，符合需求。
• 对包含多个.的字符串（如a.b.c），会正确提取最后一个.后的部分作为后缀，统计所有以.c结尾的唯一字符串数量，完全匹配endswith('.c')的逻辑。

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