如何为Pandas DataFrame创建随机80/20分区并添加buildFlag字段

可以通过以下两种方式实现需求，保证buildFlag字段随机分配且80%为1、20%为0：

方法一：精准控制比例的手动生成法 #

先按比例生成固定数量的1和0，再打乱顺序保证随机性：

import pandas as pd
import numpy as np

# 原DataFrame
ds = {'col1':[1.0,2.1,2.2,3.1,41,5.2,5.0,6.1,7.1,10]}
df = pd.DataFrame(data=ds)

# 计算80%数据的行数
total_rows = len(df)
train_count = int(total_rows * 0.8)

# 生成对应比例的标识数组并打乱
build_flag = np.concatenate([np.ones(train_count), np.zeros(total_rows - train_count)])
np.random.shuffle(build_flag)

# 添加到原DataFrame并转为整数类型
df['buildFlag'] = build_flag.astype(int)

print(df)

方法二：简洁的概率抽样法 #

直接利用numpy.random.choice指定概率生成随机序列：

import pandas as pd
import numpy as np

# 原DataFrame
ds = {'col1':[1.0,2.1,2.2,3.1,41,5.2,5.0,6.1,7.1,10]}
df = pd.DataFrame(data=ds)

# 按80%/20%概率生成buildFlag
df['buildFlag'] = np.random.choice([1, 0], size=len(df), p=[0.8, 0.2]).astype(int)

print(df)

补充说明 #

• 如果需要固定随机结果（可复现），可以在代码开头添加np.random.seed(123)（数字可自定义），锁定随机种子。
• 两种方法都能满足需求，方法一更精准控制比例（不会因随机波动出现比例偏差），方法二更简洁。

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