本文深入探讨了在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对象和额外的内存操作。
为了解决上述问题，我们可以利用Pandas的 groupby().apply() 方法，结合一个预先构建的字典来高效地传递每个组的抽样参数。apply() 方法虽然在某些情况下不如完全向量化的操作快，但它将对每个组的操作封装在一个函数中，并由Pandas在内部进行管理，通常比显式的Python循环效率更高。
核心思路是：
首先，我们需要一个DataFrame来定义每个分组的样本数量。
import pandas as pd
import 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)输出:
df_counts: a count 0 1 1 1 2 3 2 3 2 df_original: a x 0 1 a 1 1 b 2 1 c 3 2 d 4 2 e 5 3 f 6 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 x 0 1 a 3 2 d 4 2 e 4 2 e 5 3 f 6 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 等于组大小，实际效果是抽样了所有唯一元素。
本文介绍了一种在Pandas中对大型数据集进行高效分组抽样的专业方法。通过将每个组的样本数量预先存储在一个字典中，并结合 groupby().apply() 方法和一个动态控制 replace 参数的自定义函数，我们能够克服传统方法的局限性，实现灵活且高性能的抽样。这种模式在处理复杂分组抽样需求时，是提升代码效率和可维护性的关键策略。
以上就是Pandas高效分组抽样：动态样本量与替换策略的详细内容，更多请关注php中文网其它相关文章！
每个人都需要一台速度更快、更稳定的 PC。随着时间的推移，垃圾文件、旧注册表数据和不必要的后台进程会占用资源并降低性能。幸运的是，许多工具可以让 Windows 保持平稳运行。
Copyright 2014-2025 https://www.php.cn/ All Rights Reserved | php.cn | 湘ICP备2023035733号
625
收藏
