如何以Pythonic方式重构代码获取Pandas DataFrame高频元素

需求说明 #

现有如下结构的Pandas DataFrame，objects列存储单张图像中出现的物体列表：

原代码通过iterrows逐行遍历，筛选出单图不同物体数量≤3的行，收集这些行中的所有不同物体并统计总频率，但iterrows在数据量较大时性能极低，需要重构代码避免使用它。

原代码（存在性能问题） #

result = []

# Iterate through rows of the dataframe
for i, row in df.iterrows():
    # Count the frequency of each object in the image
    frequencies = Counter(row['objects'])
    # Sort the frequencies from most to least common
    sorted_frequencies = sorted(frequencies.items(),
                                    key=lambda x: x[1],
                                    reverse=True
                                    )

    # Check if there are less than 3 different objects in the image
    if len(sorted_frequencies) <= 3:
        # If so, append all of the objects to the result list
        result.extend([obj for obj, _ in sorted_frequencies])

frequency_3_most_pop = dict(Counter(result))

优化方案 #

方案1：使用apply替代iterrows #

这种方式比iterrows高效，减少逐行遍历的开销：

from collections import Counter
from itertools import chain

# 计算每行的唯一物体数量
df['unique_obj_count'] = df['objects'].apply(lambda x: len(Counter(x)))

# 筛选符合条件的行
filtered_rows = df[df['unique_obj_count'] <= 3]['objects']

# 提取所有符合条件行的唯一物体并展平
all_target_objects = list(chain.from_iterable(filtered_rows.apply(lambda x: list(Counter(x).keys()))))

# 统计总频率
frequency_3_most_pop = dict(Counter(all_target_objects))

方案2：全矢量化操作（性能最优） #

利用Pandas内置的explode和分组统计，完全避免Python层面的循环，大数据量下性能提升明显：

# 展开objects列，保留原行索引
exploded_df = df.explode('objects').reset_index(names='original_index')

# 按原索引分组，统计每行的唯一物体数量
unique_obj_per_row = exploded_df.groupby('original_index')['objects'].nunique()

# 筛选出符合条件的原行索引
valid_indices = unique_obj_per_row[unique_obj_per_row <= 3].index

# 提取符合条件行的所有唯一物体
valid_objects = exploded_df[exploded_df['original_index'].isin(valid_indices)]\
    .drop_duplicates(subset=['original_index', 'objects'])\
    ['objects']

# 统计频率（直接用Pandas的value_counts，比Counter更高效）
frequency_3_most_pop = valid_objects.value_counts().to_dict()

说明 #

两种优化方案都和原代码逻辑完全一致：仅收集单图不同物体数≤3的行中的不同物体，再统计这些物体在所有符合条件行中出现的次数。其中方案2的矢量化操作性能最优，推荐在数据量较大时使用。

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