Python中随机删除DataFrame行以平衡多任务数据集数量

嘿，这个场景太常见了——处理类别不平衡的下采样嘛，用Pandas几行代码就能解决，我给你一步步拆解：

核心思路 #

咱们要做的就是对数量过多的t1类别进行随机下采样，把它的样本量降到和t2一致，同时保留t3、t4、t5的所有数据。

具体代码实现 #

假设你的DataFrame名叫df，按下面的步骤来：

1. 先获取t2类别的样本数量：

import pandas as pd

# 获取t2的样本数
t2_sample_count = df[df['task'] == 't2'].shape[0]

2. 从t1的所有行里随机抽取和t2数量相同的样本：

# 随机下采样t1，random_state固定种子保证结果可复现
t1_downsampled = df[df['task'] == 't1'].sample(n=t2_sample_count, random_state=42)

要是不需要固定随机结果，去掉random_state=42就行，每次运行会得到不同的随机样本。

3. 合并下采样后的t1、原t2以及其他任务的数据：

# 获取t3、t4、t5的所有数据
other_tasks_data = df[df['task'].isin(['t3', 't4', 't5'])]

# 合并所有数据并重置索引
balanced_df = pd.concat([t1_downsampled, df[df['task'] == 't2'], other_tasks_data]).reset_index(drop=True)

验证结果 #

你可以用下面的代码检查处理后各任务的样本量是否符合预期：

print(balanced_df['task'].value_counts())

输出里t1和t2的数量应该完全一致，t3-t5保持原有数量。

扩展方案（通用下采样） #

如果你之后需要把所有任务的样本量都统一到当前最小的那个类别数量，可以用更通用的分组采样方法：

# 获取所有任务中最小的样本量
min_sample_count = df['task'].value_counts().min()

# 对每个任务都采样min_sample_count条数据
balanced_df = df.groupby('task').apply(lambda x: x.sample(n=min_sample_count, random_state=42)).reset_index(drop=True)

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