内存高效实现DataFrame按key聚合：分类列归一化计数+数值列均值

针对大型DataFrame因explode展开导致内存溢出的问题，我们可以直接基于原始数据的cnt列做加权聚合，完全不需要展开数据。核心思路是利用cnt作为权重，直接计算分类列的归一化频次和数值列的加权均值。

步骤分解 #

• 先计算每个key对应的总重复次数（后续计算占比和均值都会用到）
• 对分类列（如X1）：按key+分类值分组求和cnt，再除以对应key的总次数得到归一化占比
• 对数值列（如X2）：计算加权均值，即(数值*cnt)的总和除以key的总次数
• 合并两类结果得到最终聚合结果

代码实现 #

用你提供的示例DataFrame演示：

import pandas as pd

# 示例数据
df = pd.DataFrame({
    'key': [1,1,1,2,2,3],
    'cnt': [8,3,4,2,6,3],
    'X1': ['a','b','a','b','a','a'],
    'X2': [1,0,0,1,0,1]
})

# 1. 计算每个key的总cnt
total_cnt = df.groupby('key')['cnt'].sum()

# 2. 处理分类列X1的归一化value_counts
# 先按key+X1分组求和cnt
x1_counts = df.groupby(['key', 'X1'])['cnt'].sum()
# 归一化：除以对应key的总cnt
x1_normalized = x1_counts.div(total_cnt, level='key').unstack(fill_value=0)
# 重命名列，更清晰
x1_normalized.columns = [f'X1_{col}' for col in x1_normalized.columns]

# 3. 处理数值列X2的加权均值
x2_mean = (df['X2'] * df['cnt']).groupby(df['key']).sum() / total_cnt
x2_mean = x2_mean.rename('X2_mean')

# 4. 合并结果
final_result = pd.concat([x1_normalized, x2_mean], axis=1)
print(final_result)

输出结果 #

X1_a  X1_b  X2_mean
key                      
1    0.727  0.273    0.615
2    0.750  0.250    0.250
3    1.000  0.000    1.000

内存高效的原因 #

• 全程没有对数据做explode展开，所有计算都基于原始DataFrame的行数，内存占用仅和原始数据规模相关
• 分组聚合操作是Pandas内部优化过的高效计算，避免了生成中间大列表/大DataFrame的内存开销

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