如何对Pandas DataFrame按条件抽样：限制个体观测数上限

你遇到的这个问题其实用Pandas的groupby方法就能轻松搞定，完全不用手动循环构建字典，还能彻底避开索引的麻烦，而且效率更高，特别适配你那2万行的真实数据场景。

下面给你两种适配不同需求的解决方案：

1. 保留原顺序，取每个个体的前N行 #

如果只是想简单截断每个个体的观测数到设定最大值（比如你示例里的2行），直接用groupby().head()就可以：

import pandas as pd

# 模拟你的示例数据
data = {
    '个体ID': ['A', 'A', 'A', 'B', 'B', 'C'],
    '数值1': [1, 2, 3, 4, 5, 6],
    '数值2': [7, 8, 9, 10, 11, 12]
}
df = pd.DataFrame(data)

max_rows_per_ind = 2
# 按个体ID分组，每个组取前2行
sampled_df = df.groupby('个体ID').head(max_rows_per_ind)

# 要是需要重置索引（避免原索引混乱），加上这一行
sampled_df = sampled_df.reset_index(drop=True)

这个方法会自动保留原数据的顺序，分组合并后也不会有索引冲突，reset_index(drop=True)能帮你生成连续的新索引，彻底解决索引困扰。

2. 随机抽样每个个体的N行 #

如果你的需求是随机抽样（而非固定取前几行），可以用groupby().sample()方法，它会自动处理组内行数不足最大值的情况（直接取该组全部数据）：

max_rows_per_ind = 2
# 随机抽样，random_state保证结果可复现
sampled_df = df.groupby('个体ID').sample(
    n=max_rows_per_ind,
    replace=False,  # 不重复抽样，设为True则允许重复抽取同一行
    random_state=42
).reset_index(drop=True)

为什么这个方法比循环好？ #

• 效率更高：Pandas的groupby是矢量化操作，比手动循环快得多，对于你2万行、400个个体的数据，性能提升会很明显；
• 自动处理索引：不用手动拼接字典，分组后的合并会自动处理索引，加上reset_index就能得到干净的连续索引；
• 代码更简洁：一行核心代码就能完成需求，可读性和维护性都比循环好太多。

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