如何按分类列指定行数拆分Pandas DataFrame？

针对你需要按Animal列分组，拆分出每组固定数量数据的需求，这里提供两种实用方法：

先模拟示例数据（方便测试） #

先创建一个和你补充示例匹配的小数据集：

import pandas as pd
import numpy as np

# 生成cat和dog各5条数据
animals = ['cat']*5 + ['dog']*5
values = np.random.randint(1, 100, size=10)
df = pd.DataFrame({'Animal': animals, 'value': values})

方法1：按原数据顺序取固定条数 #

如果不需要随机，直接取每个动物组的前N条：

# 每个动物取3条（对应你的示例需求）
df_part1 = df.groupby('Animal').head(3)
# 剩下的2条自动分到第二个DataFrame
df_part2 = df.drop(df_part1.index)

验证结果：

print(df_part1['Animal'].value_counts())
# 输出：
# cat    3
# dog    3

print(df_part2['Animal'].value_counts())
# 输出：
# cat    2
# dog    2

方法2：随机抽取固定条数（避免顺序偏差） #

如果需要随机拆分每组数据，用sample方法：

# 每个动物随机抽3条，random_state固定随机种子保证结果可复现
df_part1 = df.groupby('Animal').sample(n=3, random_state=42)
# 剩余数据
df_part2 = df.drop(df_part1.index)

适配你的实际需求 #

把上述代码中的3换成40，2换成10即可：

# 每个动物取40条
df_main = df.groupby('Animal').head(40)  # 或用sample(n=40, random_state=42)实现随机抽取
# 剩余的10条
df_remain = df.drop(df_main.index)

关键逻辑说明 #

• groupby('Animal')：将DataFrame按Animal的唯一值分组，每个组对应一种动物的所有数据
• head(n)：保留每个组的前n条数据，保持原数据顺序
• sample(n)：从每个组随机抽取n条数据，适合需要无偏拆分的场景
• drop(index)：通过已抽取数据的索引，从原DataFrame中移除，得到剩余部分

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