如何为Pandas DataFrame生成符合指定规则的随机值列col2

问题说明 #

已通过以下代码创建DataFrame：

import pandas as pd
import numpy as np

ds = {'col1':[1,1,1,1,1,0,0,0]}
df = pd.DataFrame(data=ds)

输出结果：

col1
0     1
1     1
2     1
3     1
4     1
5     0
6     0
7     0

需要新增列col2，满足：

• 当col1=1时，5条记录中随机3条赋值为2，剩余2条赋值为3；
• 当col1=0时，3条记录中随机2条赋值为5，剩余1条赋值为6。

解决方案 #

可以通过groupby结合随机打乱的方式实现，以下是两种简洁的实现方式：

方法一：自定义分组处理函数 #

import pandas as pd
import numpy as np

ds = {'col1':[1,1,1,1,1,0,0,0]}
df = pd.DataFrame(data=ds)

def generate_col_vals(group):
    if group.name == 1:
        # 生成3个2和2个3的列表
        vals = [2]*3 + [3]*2
    else:
        # 生成2个5和1个6的列表
        vals = [5]*2 + [6]
    # 随机打乱列表顺序
    np.random.shuffle(vals)
    # 返回与分组索引匹配的序列
    return pd.Series(vals, index=group.index)

# 按col1分组处理，生成col2
df['col2'] = df.groupby('col1', group_keys=False).apply(generate_col_vals)

print(df)

方法二：使用transform简化写法 #

import pandas as pd
import numpy as np

ds = {'col1':[1,1,1,1,1,0,0,0]}
df = pd.DataFrame(data=ds)

# 利用transform直接生成对应分组的随机值序列
df['col2'] = df.groupby('col1')['col1'].transform(
    lambda x: np.random.permutation([2]*3 + [3]*2) if x.name == 1 
              else np.random.permutation([5]*2 + [6])
)

print(df)

说明 #

• 由于使用了随机打乱操作，每次运行代码生成的col2结果都会不同，符合需求中的随机位置要求；
• 两种方法均能保证每个分组内的值数量完全符合要求，不会出现数量偏差。

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