Pandas大数据集分组抽样：灵活控制每组样本量与替换行为

本文详细介绍了如何在pandas中对大型数据集进行高效的分组抽样，特别是当每个分组需要不同数量的样本，并且需要根据分组大小动态调整是否允许重复抽样（`replace`参数）时。通过构建一个自定义函数并结合`groupby().apply()`方法，可以实现比传统循环更优的性能和更灵活的控制，确保抽样逻辑的准确性和效率。

在数据分析和机器学习任务中，我们经常需要从数据集中抽取样本。当数据集非常庞大，并且需要根据某个或多个列进行分组，然后对每个分组执行抽样时，情况会变得复杂。更进一步，如果每个分组所需的样本数量不同，并且需要根据分组的实际大小动态决定是否允许重复抽样，那么传统的groupby().sample()方法就显得力不从心。

问题背景：传统方法的局限性

Pandas提供了DataFrame.groupby().sample(n=...)方法，可以方便地对每个分组抽取固定数量的样本。例如：

df.groupby("a").sample(n=1, random_state=1)

然而，这个方法无法直接处理每个分组需要不同n值的情况。一种常见的、但效率低下的做法是遍历每个分组，然后单独进行过滤和抽样：

# 假设 'm_values' 是一个字典或列表，包含每个组 'a' 对应的样本数量
for group_val, m in m_values.items():
    filter_df = df.loc[(df['a'] == group_val)]
    # ... 进行抽样，并根据条件设置 replace 参数
    # ... 最后将结果合并

这种基于循环的方法在处理拥有大量唯一分组（例如10万个）的大型数据集（例如9000万行）时，会因为频繁的数据过滤、创建子DataFrame以及最终的合并操作而导致性能瓶颈。

此外，关于replace参数的动态设置也是一个关键点：

• 如果分组的记录数小于所需样本数n，为了达到n个样本，必须允许重复抽样（replace=True）。
• 如果分组的记录数大于或等于所需样本数n，通常我们希望抽取n个唯一的样本（replace=False）。

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

Pandas的groupby().apply()方法允许我们将一个自定义函数应用到每个分组上，从而实现复杂的、组级别的数据操作。这是解决上述问题的理想方案，因为它能将Python循环的逻辑“推”到C层面执行，从而获得更好的性能。

1. 准备样本数量映射

首先，我们需要一个机制来告诉每个分组应该抽取多少样本。这通常可以通过一个包含分组键和对应样本数量的DataFrame来构建，并将其转换为字典，以便快速查找。

假设我们有一个DataFrame df1，其中包含每个组a所需的样本数量：

壁纸样机神器

免费壁纸样机生成

0

查看详情

import pandas as pd

# df1: 定义每个组 'a' 对应的样本数量
df1 = pd.DataFrame({
    'a': [1, 2, 3],
    'count': [1, 3, 2]
})

# 将 df1 转换为字典，方便查找
sample_counts_map = df1.set_index("a")["count"].to_dict()
# sample_counts_map -> {1: 1, 2: 3, 3: 2}

2. 定义自定义抽样函数

接下来，我们创建一个函数，该函数将作为apply()的参数，对每个分组DataFrame进行操作。这个函数需要接收分组DataFrame、样本数量映射字典以及随机种子作为参数。

def get_sample(group_df, sample_counts_dict, random_state):
    """
    对每个分组DataFrame进行抽样。
    根据分组键从 sample_counts_dict 获取所需的样本数量，
    并根据分组大小动态调整 replace 参数。

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

    Returns:
        pd.DataFrame: 抽样后的DataFrame，如果该组不需要抽样则返回 None。
    """
    # 获取当前分组的键值 (例如 'a' 列的值)
    group_key = group_df["a"].iat[0] # iat[0] 效率更高，因为我们知道组内 'a' 值都相同

    # 从映射字典中获取该组所需的样本数量
    n_samples = sample_counts_dict.get(group_key)

    # 如果该组的键不在映射字典中，或者 n_samples 为 None，则不进行抽样
    if n_samples is None:
        return None

    # 动态设置 replace 参数
    # 如果分组的实际记录数小于或等于所需样本数 n_samples，则必须允许重复抽样 (replace=True)
    # 否则，如果分组记录数大于 n_samples，则不允许重复抽样 (replace=False)
    allow_replace = len(group_df) <= n_samples

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

3. 应用 groupby().apply()

最后，我们将这个自定义函数应用到原始DataFrame的groupby对象上。

# df2: 原始数据，待抽样
df2 = pd.DataFrame({
    'a': [1, 1, 1, 2, 2, 3, 3],
    'x': ['a', 'b', 'c', 'd', 'e', 'f', 'g']
})

# 使用 groupby().apply() 进行抽样
# group_keys=False 可以避免在结果中添加分组键作为索引，提高性能
sampled_df = df2.groupby("a", group_keys=False).apply(
    get_sample,
    sample_counts_dict=sample_counts_map,
    random_state=6
)

print(sampled_df)

输出结果：

   a  x
0  1  a
3  2  d
4  2  e
4  2  e
5  3  f
6  3  g

从结果可以看出：

• 组 a=1：原始有3条记录，需要1条样本，结果是1条（replace=False）。
• 组 a=2：原始有2条记录，需要3条样本，结果是3条，其中一条是重复的（replace=True）。
• 组 a=3：原始有2条记录，需要2条样本，结果是2条（replace=False）。

这完美地实现了我们动态抽样的需求。

关键点与注意事项

1. group_keys=False: 在groupby().apply()中设置group_keys=False是一个重要的性能优化。它阻止了apply操作在结果DataFrame中将分组键作为额外的索引层，这在处理大量分组时可以显著减少内存开销和处理时间。
2. random_state: 设置random_state参数可以确保每次运行代码时，抽样结果都是一致的，这对于调试和结果复现至关重要。
3. sample_counts_dict 的效率: 将样本数量映射预先构建成字典，使得在自定义函数中通过dict.get()查找样本数量的操作非常高效，是O(1)的平均时间复杂度。
4. len(group_df) <= n_samples 逻辑: 理解这个条件对于正确实现动态replace行为至关重要。当分组的记录数不足以提供所需样本数时，必须允许重复抽样；否则，为了获取唯一样本，应禁止重复抽样。
5. 内存管理: 尽管apply()比显式循环更高效，但对于极大的分组，group_df本身可能仍然占用大量内存。在处理数亿甚至数十亿行的数据集时，可能需要考虑使用更高级的分布式计算框架（如PySpark、Dask）来避免单机内存限制。

总结

通过利用Pandas的groupby().apply()方法结合一个精心设计的自定义抽样函数，我们可以高效地解决大型数据集上复杂的分组抽样问题。这种方法不仅能够灵活地为每个分组指定不同的样本数量，还能根据分组的实际大小智能地调整是否允许重复抽样，从而在保证数据质量的同时，显著提升处理效率。

以上就是Pandas大数据集分组抽样：灵活控制每组样本量与替换行为的详细内容，更多请关注php中文网其它相关文章！