Python中如何无循环按ID分组生成和为1的多项分布target列？

现有如下DataFrame：

id  test
0   1   1.000000
1   2   0.582594
2   2   0.417406
3   3   0.016633
4   3   0.983367
5   4   1.000000
6   5   0.501855
7   5   0.498145
8   6   1.000000
9   7   1.000000

需要使用np.random.multinomial()函数生成新列target，要求：

• 每个ID对应的target值根据该ID下test列的概率作为pvals参数生成0或1
• 每个ID的target列总和恒为1

例如针对id=2，执行np.random.multinomial(n = 1, pvals = [0.582594, 0.417406])可能得到array([1, 0])，最终期望结果类似：

id  test       target
0   1   1.000000   1
1   2   0.582594   1
2   2   0.417406   0
3   3   0.016633   0
4   3   0.983367   1
5   4   1.000000   1 
6   5   0.501855   0
7   5   0.498145   1
8   6   1.000000   1
9   7   1.000000   1

要求不通过循环迭代每个ID实现该需求。

可以利用Pandas的分组功能结合Numpy的向量化操作实现，无需显式循环：

1. 按ID分组提取概率数组
   对DataFrame按id分组，提取每组的test值作为概率列表，同时保留组内的行对应关系。

2. 批量生成多项分布结果
   对每个分组的概率数组调用np.random.multinomial()，设置n=1确保每组结果只有一个1，其余为0，满足总和为1的要求。

3. 展开结果并映射回原DataFrame
   将分组生成的数组拆分为单独行，直接匹配原DataFrame的行顺序，赋值为新列。

具体代码如下：

import pandas as pd
import numpy as np

# 构造原始DataFrame
df = pd.DataFrame({
    'id': [1,2,2,3,3,4,5,5,6,7],
    'test': [1.0, 0.582594, 0.417406, 0.016633, 0.983367, 1.0, 0.501855, 0.498145, 1.0, 1.0]
})

# 按id分组生成目标列数据
target = (
    df.groupby('id')['test']
      .apply(lambda x: np.random.multinomial(n=1, pvals=x.values))
      .explode()
      .astype(int)
)

# 将结果赋值回原DataFrame
df['target'] = target.values

print(df)

代码说明 #

• groupby('id')['test']：按ID分组，聚焦每组的概率值
• apply(lambda x: np.random.multinomial(n=1, pvals=x.values))：对每组概率生成多项分布结果，每组得到一个长度与组内行数一致的数组，仅含一个1
• explode()：将每组的数组拆分为独立行，与原DataFrame的行一一对应
• astype(int)：将结果转为整数类型（0或1）
• 最后将展开后的结果赋值给df['target']，完成目标列生成

这种方法借助Pandas分组的内置优化，避免手动循环，处理大数据量时效率更高。

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