如何在Pandas分组数据中跨行跨列比较值并过滤分组

问题说明 #

现有10000×61的数据集，需按ColumnA和ColumnB分组，保留同时满足以下条件的分组，不满足则删除整个分组：

• 分组内每一行的ColumnD与ColumnC的天数差>900
• 分组内除最后一行外的每一行，下一行的ColumnC与当前行ColumnD的天数差<70
• 需处理ColumnD的空值，避免日期计算错误

解决方案 #

核心思路 #

1. 将日期列转换为Pandas datetime类型，空值自动转为NaT（Pandas缺失日期标记），避免引入无关日期干扰判断
2. 直接使用groupby(['ColumnA', 'ColumnB'])进行分组，无需手动生成分组ID
3. 编写分组过滤函数，对每个分组逐一验证条件，仅保留符合要求的分组

完整代码 #

import pandas as pd

def filter_group(group):
    # 分组至少需要2行（否则无下一行可比较）
    if len(group) < 2:
        return False
    
    # 条件2：所有行的ColumnD - ColumnC 天数差 > 900
    cond2 = (group['ColumnD'] - group['ColumnC']).dt.days > 900
    # 空值会导致cond2为False，只要有一行不满足则排除该分组
    if not cond2.all():
        return False
    
    # 条件1：下一行ColumnC - 当前行ColumnD 天数差 <70（仅检查前n-1行）
    cond1 = (group['ColumnC'].shift(-1) - group['ColumnD']).dt.days[:-1] < 70
    # 空值会导致cond1为False，只要有一行不满足则排除该分组
    if not cond1.all():
        return False
    
    return True

def main():
    # 示例数据
    data = pd.DataFrame([
        ['Value1', 'Value2','2010-12-06' , '2014-03-14'],
        ['Value1', 'Value2','2014-04-13', '2017-10-01'], 
        ['Value1', 'Value2','2017-11-01',''],
        ['Value3', 'Value4', '2010-12-06', '2011-03-14'], 
        ['Value3', 'Value4', '2014-04-13', '2017-10-01']
    ], columns=['ColumnA', 'ColumnB', 'ColumnC', 'ColumnD'])
    
    # 转换日期列，空值自动转为NaT
    data['ColumnC'] = pd.to_datetime(data['ColumnC'], format='%Y-%m-%d')
    data['ColumnD'] = pd.to_datetime(data['ColumnD'], format='%Y-%m-%d')
    
    # 按ColumnA和ColumnB分组并过滤
    filtered_data = data.groupby(['ColumnA', 'ColumnB']).filter(filter_group)
    print(filtered_data)

if __name__ == "__main__":
    main()

代码解释 #

1. 日期转换：使用pd.to_datetime将ColumnC和ColumnD转为datetime类型，ColumnD的空值会自动转为NaT，避免手动填充无关日期导致的判断误差
2. 分组过滤函数：
   • 先判断分组行数，少于2行直接排除（无下一行可比较条件1）
   • 条件2验证：计算每行ColumnD - ColumnC的天数差，检查是否所有行都大于900；NaT对应的天数差为NaN，会被视为False，因此包含空值的分组会被排除
   • 条件1验证：使用shift(-1)获取下一行的ColumnC，计算与当前行ColumnD的天数差，仅检查前n-1行（最后一行无下一行）；同理，空值会导致条件不满足，分组被排除
3. 执行过滤：通过groupby().filter()直接保留符合条件的分组

内容的提问来源于stack exchange，提问作者pyflo