如何不展开聚合数据为Pandas分类统计结果快速分配概率标签

最优解决方案 #

核心思路是利用多项分布直接抽样计数，完全不需要展开原数据的聚合行，时间复杂度仅和原DataFrame的行数正相关，和每行的number数值大小无关，相比暴力展开方案性能提升可达几个数量级，完全适配百万级行的业务场景。

代码实现

import pandas as pd
import numpy as np

# 示例数据准备
df = pd.DataFrame(data={
    'color': ['red', 'red', 'green', 'green', 'blue', 'blue'], 
    'shape': ['round', 'square', 'round', 'square', 'round', 'square'], 
    'number': [10000, 1, 50, 500, 7, 3000]
})
score_dist = {'A': 0.1, 'B': 0.5, 'C': 0.25, 'D': 0.15}

# 提取标签和对应概率
score_labels = list(score_dist.keys())
score_probs = list(score_dist.values())

# 按行生成各标签的抽样计数
def generate_label_counts(row):
    # 多项分布直接抽样，参数为总样本数、各标签概率
    counts = np.random.multinomial(row['number'], score_probs)
    return pd.Series(counts, index=score_labels)

# 合并结果并转换为目标格式
counts_matrix = df.apply(generate_label_counts, axis=1)
res = pd.concat([df[['color', 'shape']], counts_matrix], axis=1)
# 宽表转长表，过滤计数为0的条目，排序得到最终结果
res = res.melt(
    id_vars=['color', 'shape'], 
    var_name='scores', 
    value_name='number'
).query('number > 0').sort_values(
    by=['color', 'shape', 'scores']
).reset_index(drop=True)

方案优势

1. 性能极高：即使原df有百万行，也能在秒级完成计算，不会出现内存溢出问题
2. 结果和暴力展开逻辑完全等价：都是服从给定概率分布的随机抽样结果
3. 代码简洁易维护：不需要处理展开、合并、逐行匹配标签等冗余逻辑

注意事项

• 可提前校验sum(score_probs) ≈ 1，避免出现概率分布不合法的警告
• 如果存在number=0的行，可在generate_label_counts函数开头加判断直接返回全0序列，规避异常

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