Pandas高效筛选：基于分组条件提取DataFrame子集

本文详细介绍了如何在pandas中，不创建额外辅助列的情况下，根据`groupby`分组聚合的条件来筛选dataframe的子集。核心方法是利用`groupby().transform()`函数，它能将分组计算结果广播回原dataframe的索引，从而实现高效的布尔索引过滤，避免了冗余数据和复杂的合并操作。

在数据分析和处理中，我们经常需要根据复杂的条件来筛选DataFrame中的行。其中一种常见场景是，筛选出那些其所属分组不满足特定聚合条件的行。例如，在一个包含日期、地点和代理人信息的数据集中，我们可能需要找出那些在特定日期和地点，代理人数量超过预设阈值的全部记录。传统方法可能涉及先计算分组聚合值，然后将其合并回原DataFrame，再进行筛选，但这会引入额外的列，有时并不高效。

1. 问题场景与数据准备

假设我们有如下一个DataFrame，记录了不同日期和地点的代理人信息：

import pandas as pd
import io

data = """date|point|agent
2023-10-02|A|agent1
2023-10-02|A|agent2
2023-10-05|B|agent3
2023-10-05|B|agent2
2023-10-02|C|agent1
2023-10-02|C|agent2
2023-10-02|C|agent3
"""
df = pd.read_csv(io.StringIO(data), sep='|')
print("原始DataFrame:")
print(df)

我们的目标是：找出在任何给定date和point组合下，agent数量超过2的那些所有行。

2. 理解 groupby().transform() 的作用

Pandas的groupby()操作通常会伴随着聚合函数（如sum(), count(), mean()等），这些函数会将每个组的数据压缩成一个单一的值，并返回一个行数少于原始DataFrame的新Series或DataFrame。

然而，transform()方法则不同。它在每个组上应用一个函数，但返回一个与原始DataFrame具有相同索引和行数的Series（或DataFrame）。这意味着transform()的结果可以被直接用于原始DataFrame的布尔索引，而无需进行合并操作。

对于本例，我们需要计算每个('point', 'date')组合中agent的唯一数量。transform('nunique')将为原始DataFrame中的每一行返回其所属分组的唯一代理人数量。

腾讯交互翻译

腾讯AI Lab发布的一款AI辅助翻译产品

183

查看详情

3. 高效筛选解决方案

利用groupby().transform('nunique')，我们可以一步到位地生成一个布尔Series，用于直接筛选DataFrame。

# 1. 计算每个(point, date)分组中agent的唯一数量，并将结果广播回原始DataFrame的索引
agent_counts_per_group = df.groupby(['point', 'date'])['agent'].transform('nunique')

# 2. 基于条件筛选出唯一代理人数量大于2的行
filtered_df = df[agent_counts_per_group > 2]

print("\n筛选结果DataFrame:")
print(filtered_df)

输出结果：

原始DataFrame:
         date point   agent
0  2023-10-02     A  agent1
1  2023-10-02     A  agent2
2  2023-10-05     B  agent3
3  2023-10-05     B  agent2
4  2023-10-02     C  agent1
5  2023-10-02     C  agent2
6  2023-10-02     C  agent3

筛选结果DataFrame:
         date point   agent
4  2023-10-02     C  agent1
5  2023-10-02     C  agent2
6  2023-10-02     C  agent3

正如所见，('2023-10-02', 'C')这个分组有三个代理人（agent1, agent2, agent3），因此所有属于这个分组的行都被成功筛选出来。

4. 总结与注意事项

• 简洁高效： 这种方法避免了创建临时列，使得代码更加简洁，内存占用更低，尤其适用于大型数据集。
• 适用场景： transform()非常适合需要将分组聚合结果“映射”回原始DataFrame的每个元素，以进行进一步筛选或计算的场景。
• 与其他方法的对比：
  • groupby().agg()会返回一个聚合后的DataFrame，其索引是分组键，无法直接用于原始DataFrame的布尔索引。
  • groupby().apply()可以执行更复杂的自定义函数，但通常比transform()效率低，因为它需要将每个组作为独立的DataFrame传递给函数。
• 性能考量： 对于非常大的数据集，transform()通常是高度优化的，但在某些极端复杂的自定义转换场景下，可能需要评估其性能。

通过熟练运用groupby().transform()，Pandas用户可以更优雅、高效地处理基于分组条件的DataFrame筛选任务，提升数据处理的灵活性和效率。

以上就是Pandas高效筛选：基于分组条件提取DataFrame子集的详细内容，更多请关注php中文网其它相关文章！