本文深入探讨了在pandas中对大型数据集进行分组抽样的优化策略，特别是当每个组需要不同的样本数量，并根据组内元素总数动态调整是否允许重复抽样（`replace`参数）时。通过介绍一种结合字典映射和`groupby().apply()`方法的解决方案，我们展示了如何避免低效的循环操作，从而显著提升处理效率，实现灵活且高性能的数据抽样。
在数据分析和机器学习任务中，我们经常需要从大型数据集中进行抽样。当数据需要根据某个或多个列进行分组，并且每个组的抽样需求（例如，抽样数量 n 和是否允许重复 replace）都不同时，传统的 df.groupby().sample() 方法可能无法满足所有要求。尤其对于包含数千万甚至上亿条记录的数据集，以及数十万个唯一分组的情况，低效的抽样方法会导致严重的性能瓶颈。
问题背景与传统方法的局限性
假设我们有一个大型DataFrame df，其中包含一个分组列 “a”，以及一个记录每个组所需抽样数量的DataFrame df_counts。我们的目标是根据 df_counts 中为每个 “a” 组指定的 count 值进行抽样。
1. 基础的 groupby().sample()：Pandas提供了 df.groupby(“a”).sample(n=k)，这可以对每个组抽取固定数量 k 的样本。但这种方法无法为每个组指定不同的 n 值。
# 示例：每个组抽取1个样本# df.groupby("a").sample(n=1, random_state=1)
这无法满足不同组不同 n 的需求。
2. 循环迭代的低效性：一种直观但效率低下的方法是遍历每个唯一的组，筛选出该组的数据，然后进行抽样，最后将所有结果拼接起来。此外，为了更灵活地控制抽样行为，我们可能还需要根据组内元素的总数与所需样本量 n 的关系，动态地决定 replace 参数（即当组内元素不足 n 时允许重复抽样，否则不允许）。
# 伪代码：循环迭代的低效方法# sampled_dfs = []# for group_val in df['a'].unique():# filter_df = df.loc[df['a'] == group_val]# n_samples = get_n_from_df_counts(group_val) # 从df_counts获取该组的n值# # if len(filter_df) >= n_samples:# sampled_group = filter_df.sample(n=n_samples, random_state=6, replace=False)# else:# sampled_group = filter_df.sample(n=n_samples, random_state=6, replace=True)# sampled_dfs.append(sampled_group)## final_sampled_df = pd.concat(sampled_dfs)
对于拥有100k个唯一值的分组列和90M行的数据集，这种基于Python循环的逐组筛选和抽样会带来巨大的性能开销，因为每次迭代都会产生新的DataFrame对象和额外的内存操作。
优化方案：结合 groupby().apply() 和字典映射
为了解决上述问题，我们可以利用Pandas的 groupby().apply() 方法，结合一个预先构建的字典来高效地传递每个组的抽样参数。apply() 方法虽然在某些情况下不如完全向量化的操作快，但它将对每个组的操作封装在一个函数中，并由Pandas在内部进行管理，通常比显式的Python循环效率更高。
核心思路是：
构建样本量字典： 将包含每个组所需样本量的数据转换成一个字典，其中键是分组列的值，值是对应的样本量。定义自定义抽样函数： 创建一个函数，该函数接收一个组的DataFrame，并根据字典查找该组的样本量 n。同时，根据组的实际大小与 n 的关系，动态设置 sample() 方法的 replace 参数。应用自定义函数： 使用 df.groupby().apply() 将自定义抽样函数应用到每个组。
步骤一：准备样本计数数据
首先，我们需要一个DataFrame来定义每个分组的样本数量。
import pandas as pdimport numpy as np# 示例输入数据# df_counts 定义了每个 'a' 组需要抽样的数量df_counts = pd.DataFrame({ 'a': [1, 2, 3], 'count': [1, 3, 2]})# 原始数据 df_original，我们将从中抽样df_original = pd.DataFrame({ 'a': [1, 1, 1, 2, 2, 3, 3], 'x': ['a', 'b', 'c', 'd', 'e', 'f', 'g']})print("df_counts:")print(df_counts)print("ndf_original:")print(df_original)
输出:
壁纸样机神器
免费壁纸样机生成
0 查看详情 
df_counts: a count0 1 11 2 32 3 2df_original: a x0 1 a1 1 b2 1 c3 2 d4 2 e5 3 f6 3 g
接下来，将 df_counts 转换为一个字典，以便在抽样函数中快速查找。
# 构建样本量字典sample_counts_dict = df_counts.set_index("a")["count"].to_dict()print("nsample_counts_dict:")print(sample_counts_dict)
输出:
sample_counts_dict:{1: 1, 2: 3, 3: 2}
步骤二：定义自定义抽样函数
这个函数将是 groupby().apply() 的核心。它接收一个组的DataFrame，一个包含所有组样本量的字典，以及一个随机种子。
def get_sample(group_df, sample_dict, random_state=None): """ 对给定的组DataFrame进行抽样。 参数: group_df (pd.DataFrame): 当前分组的DataFrame。 sample_dict (dict): 包含每个组所需样本量的字典。 random_state (int, optional): 随机种子，用于结果可复现性。 返回: pd.DataFrame: 抽样后的DataFrame，如果该组没有对应的样本量则返回None。 """ # 获取当前组的键（例如 'a' 列的值） group_key = group_df["a"].iat[0] # 从字典中获取该组的样本量n n_samples = sample_dict.get(group_key) # 如果字典中没有该组的样本量，则返回None（表示不抽样或跳过） if n_samples is None: return None # 动态设置 replace 参数 # 如果组内元素数量小于所需样本量n，则必须允许重复抽样 (replace=True) # 否则，如果组内元素数量足够，则默认不允许重复抽样 (replace=False) # 这样可以最大化地获取唯一行 replace_flag = len(group_df) <= n_samples return group_df.sample(n=n_samples, random_state=random_state, replace=replace_flag)
步骤三：应用自定义函数进行抽样
最后，我们将 get_sample 函数应用到 df_original 的每个组。
# 应用自定义函数进行分组抽样# group_keys=False 可以避免将分组键作为额外的索引添加到结果中，保持输出整洁sampled_output_df = df_original.groupby("a", group_keys=False).apply( get_sample, sample_dict=sample_counts_dict, random_state=6)print("nDesired sampled output:")print(sampled_output_df)
输出:
Desired sampled output: a x0 1 a3 2 d4 2 e4 2 e5 3 f6 3 g
可以看到，对于 a=1，原始数据有3条，需要抽样1条，结果是 a。对于 a=2，原始数据有2条（’d’, ‘e’），需要抽样3条。由于 len(group_df) <= n_samples (2 <= 3) 为 True，replace 被设置为 True，因此 e 被重复抽样。对于 a=3，原始数据有2条（'f', 'g'），需要抽样2条。由于 len(group_df) <= n_samples (2 <= 2) 为 True，replace 被设置为 True，但因为 n_samples 等于组大小，实际效果是抽样了所有唯一元素。
性能考量与注意事项
groupby().apply() 的优势： 相较于显式的Python循环，apply() 在内部进行了优化，尤其是在处理大量组时，可以减少Python解释器和Pandas对象之间的切换开销。它允许我们将自定义逻辑高效地应用于每个组。group_keys=False： 在 groupby().apply() 中设置 group_keys=False 可以防止分组键作为额外的索引添加到结果DataFrame中，这在结果拼接时可以避免不必要的索引重置操作，并保持输出结构的简洁。replace 参数的动态控制： 这种方法允许我们根据每个组的实际情况，灵活地设置 replace 参数，确保在组内元素不足以满足抽样数量时，能够通过重复抽样来达到目标数量，而在元素充足时，则优先进行无重复抽样。随机种子 random_state： 使用 random_state 参数可以确保抽样结果的可复现性，这在调试和结果验证时非常重要。大数据集的挑战： 尽管 apply() 比循环更优，但对于拥有极其庞大数量的组（例如数百万个组），且每个组的数据量都非常小的情况，apply() 仍然可能面临性能挑战。在这种极端情况下，可能需要考虑使用更底层的库（如Numba）或分布式计算框架（如Spark）来进一步优化。然而，对于大多数常见的大数据集场景，groupby().apply() 配合自定义函数是一个非常高效且灵活的解决方案。
总结
本文介绍了一种在Pandas中对大型数据集进行高效分组抽样的专业方法。通过将每个组的样本数量预先存储在一个字典中，并结合 groupby().apply() 方法和一个动态控制 replace 参数的自定义函数，我们能够克服传统方法的局限性，实现灵活且高性能的抽样。这种模式在处理复杂分组抽样需求时，是提升代码效率和可维护性的关键策略。
以上就是Pandas高效分组抽样：动态样本量与替换策略的详细内容，更多请关注创想鸟其它相关文章！
版权声明：本文内容由互联网用户自发贡献，该文观点仅代表作者本人。本站仅提供信息存储空间服务，不拥有所有权，不承担相关法律责任。
如发现本站有涉嫌抄袭侵权/违法违规的内容， 请发送邮件至 chuangxiangniao@163.com 举报，一经查实，本站将立刻删除。
发布者：程序猿，转转请注明出处：https://www.chuangxiangniao.com/p/592395.html
微信扫一扫 
支付宝扫一扫 
