Pandas分组抽样：动态样本与替换方法

在数据分析领域，Pandas是处理表格数据不可或缺的工具。其中，分组（groupby）操作是其核心功能之一。然而，当面临复杂的抽样需求时，例如对一个包含数千万甚至上亿数据点的大型数据集进行分组抽样，并且每个分组需要抽取不同数量的样本，同时还要根据分组的实际大小动态决定是否允许重复抽样时，传统的groupby().sample()方法可能力不从心。本文将介绍一种基于groupby().apply()的优化策略，以高效、灵活地解决此类高级分组抽样问题。

1. 挑战：动态分组抽样与替换策略

假设我们有一个庞大的数据集df2，其中包含一个分组键列'a'。我们希望根据'a'列对数据进行分组，并从每个组中抽取特定数量的样本。这些样本数量并非固定值，而是由另一个数据帧df1（例如，包含每个组'a'对应的'count'值）所定义。

更进一步的挑战在于抽样时的“替换”（replace）策略。如果一个分组的实际行数小于或等于我们希望抽取的样本数n，那么为了达到n个样本的目标，我们必须允许重复抽样（replace=True）。反之，如果分组的行数大于n，我们通常会倾向于不重复抽样（replace=False），以获得尽可能多的唯一行。

直接使用df.groupby("a").sample(n=N)无法满足每个组N不同的需求。而通过循环遍历每个组并单独抽样，如：

# 伪代码：低效的循环方法
sampled_dfs = []
for group_key in df['a'].unique():
    group_df = df[df['a'] == group_key]
    n_samples = sample_counts_dict.get(group_key) # 从预设字典获取n
    if n_samples is not None:
        if len(group_df) >= n_samples:
            sampled_group = group_df.sample(n=n_samples, random_state=6, replace=False)
        else:
            sampled_group = group_df.sample(n=n_samples, random_state=6, replace=True)
        sampled_dfs.append(sampled_group)
result = pd.concat(sampled_dfs)

这种基于循环的方法对于拥有10万个唯一分组键（如问题描述中'a'列有10万个唯一值）的大型数据集来说，性能会非常低下，因为它涉及多次数据筛选、创建子DataFrame以及拼接操作。

2. 解决方案：groupby().apply()结合自定义函数

Pandas的groupby().apply()方法提供了一种更高效、更“Pandas风格”的解决方案。它允许我们将一个自定义函数应用到groupby对象生成的每个子DataFrame上，从而在C语言级别进行优化，显著提升性能。

2.1 准备示例数据

首先，我们定义用于演示的数据帧df1（包含每个组的样本计数）和df2（待抽样的原始数据）。

import pandas as pd
import numpy as np

# df1: 定义每个组 'a' 需要抽取的样本数量
data1 = {'a': [1, 2, 3], 'count': [1, 3, 2]}
df1 = pd.DataFrame(data1)
print("df1 (样本计数):\n", df1)

# df2: 原始数据集
data2 = {'a': [1, 1, 1, 2, 2, 3, 3], 'x': ['a', 'b', 'c', 'd', 'e', 'f', 'g']}
df2 = pd.DataFrame(data2)
print("\ndf2 (原始数据):\n", df2)

输出示例:

df1 (样本计数):
    a  count
0  1      1
1  2      3
2  3      2

df2 (原始数据):
    a  x
0  1  a
1  1  b
2  1  c
3  2  d
4  2  e
5  3  f
6  3  g

2.2 构建样本计数字典

为了在自定义函数中高效地查找每个组所需的样本数量，我们将df1转换为一个字典，其中键是分组键'a'，值是对应的样本数量'count'。

sample_counts_dict = df1.set_index("a")["count"].to_dict()
print("\n样本计数字典:\n", sample_counts_dict)

输出示例:

样本计数字典:
 {1: 1, 2: 3, 3: 2}

2.3 定义自定义抽样函数

核心在于创建一个函数，它能接收一个分组DataFrame，并根据预设的字典和动态的replace逻辑进行抽样。

def get_sample_per_group(group_df, sample_counts_dict, random_state):
    """
    根据每个组的样本计数字典，对当前分组DataFrame进行抽样。
    动态调整 replace 参数：如果所需样本数大于分组大小，则 replace=True。

    Args:
        group_df (pd.DataFrame): 当前分组的DataFrame。
        sample_counts_dict (dict): 包含每个分组键及其所需样本数量的字典。
        random_state (int): 随机种子，用于保证抽样结果的可复现性。

    Returns:
        pd.DataFrame: 抽样后的DataFrame，如果该分组无需抽样则返回None。
    """
    # 获取当前分组的键值 (例如，'a' 列的值)
    # .iat[0] 比 .iloc[0] 更快，因为它直接访问单个元素
    group_key = group_df["a"].iat[0]

    # 从字典中查找当前分组所需的样本数量
    n_samples = sample_counts_dict.get(group_key)

    # 如果字典中没有对应的样本数量，则不进行抽样，返回None
    if n_samples is None:
        return None

    # 动态确定 replace 参数
    # 如果当前分组的行数小于或等于所需样本数，则必须允许重复抽样
    # 否则，默认不重复抽样以获取唯一行
    replace_flag = len(group_df) <= n_samples

    # 执行抽样
    return group_df.sample(n=n_samples, random_state=random_state, replace=replace_flag)

replace参数的动态逻辑解释：

• len(group_df)：当前分组的实际行数。
• n_samples：该分组期望抽取的样本数量。
• 当len(group_df) <= n_samples时，意味着我们希望抽取的样本数大于或等于分组的实际行数。在这种情况下，为了达到n_samples的目标，我们必须允许重复抽样，因此replace_flag为True。
• 当len(group_df) > n_samples时，意味着分组的实际行数足够多，可以从中抽取n_samples个唯一的样本。此时，replace_flag为False（或不指定，因为False是sample函数的默认值），以确保抽取到唯一的行。

2.4 应用自定义函数进行抽样

最后，我们将自定义函数get_sample_per_group应用到df2的groupby('a')对象上。

# 设置一个随机种子以确保结果可复现
RANDOM_STATE = 6

# 对 df2 按 'a' 列进行分组，并应用自定义抽样函数
# group_keys=False 可以避免在结果中将分组键作为额外的索引层
sampled_df = df2.groupby("a", group_keys=False).apply(
    get_sample_per_group,
    sample_counts_dict=sample_counts_dict,
    random_state=RANDOM_STATE
)

print("\n最终抽样结果:\n", sampled_df)

输出示例:

最终抽样结果:
    a  x
0  1  a
3  2  d
4  2  e
4  2  e
5  3  f
6  3  g

从结果可以看出：

• 对于a=1，df1中count为1，原始df2中有3行。len(group_df) (3) > n_samples (1)，所以replace=False，抽取1个样本（如a）。
• 对于a=2，df1中count为3，原始df2中有2行。len(group_df) (2) <= n_samples (3)，所以replace=True，抽取3个样本（如d, e, e，e被重复抽取）。
• 对于a=3，df1中count为2，原始df2中有2行。len(group_df) (2) <= n_samples (2)，所以replace=True，抽取2个样本（如f, g，虽然replace=True但因为数量相等，实际是全部抽取）。

3. 注意事项与总结

• 性能优势： groupby().apply()在Pandas底层进行了优化，避免了Python层面的显式循环，对于大型数据集和大量分组键，其性能远超手动循环。
• group_keys=False： 在apply操作中设置group_keys=False可以防止分组键作为额外的索引层出现在结果中，使输出更加扁平化和易于处理。
• random_state： 使用random_state参数可以确保每次运行代码时抽样结果的可复现性，这在调试和结果验证时非常重要。
• 内存管理： 尽管apply是高效的，但如果单个分组非常庞大，处理单个分组DataFrame仍然可能占用大量内存。然而，apply是逐组处理的，通常不会一次性将所有分组加载到内存。
• 错误处理： 在get_sample_per_group函数中，通过sample_counts_dict.get(group_key)并检查None值，可以优雅地处理某些分组键在df1中没有对应样本计数的情况。

通过这种groupby().apply()的组合策略，我们不仅解决了Pandas分组抽样中动态样本量和条件替换的复杂需求，而且确保了在处理大规模数据集时的性能和可扩展性。这种模式在处理各种复杂的组级操作时都非常有用，是Pandas高级应用中的一个强大工具。