Pandas多列Groupby后按组大小规则采样的技术实现问询

嘿，这个分组采样的需求我熟！刚好可以用Pandas的groupby.sample来优雅解决，完全贴合你的要求：组内元素超过2条就抽取2条样本，≤2条则保留全部记录。

步骤1：构造原始DataFrame #

先把你给出的示例数据转换成Pandas DataFrame：

import pandas as pd

data = {
    'col1': ['A1', 'A1', 'A1', 'B1', 'B1', 'C1'],
    'col2': ['A2', 'A2', 'A2', 'B2', 'B2', 'C2'],
    'col3': ['A3', 'A3', 'A3', 'B3', 'B3', 'C3'],
    'col4': ['A4', 'A5', 'A6', 'B4', 'B5', 'C4']
}
df = pd.DataFrame(data)

步骤2：核心分组采样代码 #

直接用groupby结合sample方法，通过lambda函数动态控制每个组的采样数量：

# 按col1、col2、col3三列分组，然后按规则采样
sampled_df = df.groupby(['col1', 'col2', 'col3']).sample(
    n=lambda group_len: min(group_len, 2),
    random_state=42  # 可选：固定随机种子，让采样结果可复现
)

代码解释 #

• 多列分组：groupby(['col1', 'col2', 'col3'])实现你需要的多列分组逻辑
• 动态采样数量：n=lambda group_len: min(group_len, 2)是关键——每个组传入自己的长度group_len，返回min(组长度, 2)，完美实现“超2取2，否则全留”的规则
• 随机种子：random_state=42是可选参数，加上后每次运行采样结果一致，方便调试；不需要的话可以直接删掉

示例输出 #

运行上述代码后，你会得到类似这样的结果（A组的两条样本是随机抽取的，每次可能不同）：

col1 col2 col3 col4
0   A1   A2   A3   A4
2   A1   A2   A3   A6
3   B1   B2   B3   B4
4   B1   B2   B3   B5
5   C1   C2   C3   C4

完全符合你期望的目标DataFrame格式，A组保留2条随机样本，B组和C组保留全部记录。

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