基于另一DataFrame条件筛选MultiIndex DataFrame的优化方案

需求规则 #

对MultiIndex DataFrame（df1）中的每个Subject ID，若普通DataFrame（df2）中存在该Subject ID对应的Visit_NUM ≥ df1中该Subject ID的Visit_NUM + 2，则保留df1中该Subject ID的所有行，否则删除。现有代码仅支持等于匹配，需优化为支持大于等于且实现更简洁。

数据示例 #

import pandas as pd

# df1：MultiIndex结构，索引为['Subject ID', 'Visit_NUM', 'Description']
df1 = pd.DataFrame({
    'Sex': {('002_S_0413', 0, 'DTI'): 'F',
            ('002_S_0413', 0, 'T1'): 'F',
            ('002_S_4213', 2, 'DTI'): 'F',
            ('002_S_4213', 2, 'T1'): 'F',
            ('002_S_4799', 0, 'DTI'): 'M',
            ('002_S_4799', 0, 'T1'): 'M',
            ('002_S_5178', 0, 'DTI'): 'M',
            ('002_S_5178', 0, 'T1'): 'M',
            ('002_S_5230', 2, 'DTI'): 'F',
            ('002_S_5230', 2, 'T1'): 'F'},
    'DIAGNOSIS': {('002_S_0413', 0, 'DTI'): 1.0,
                  ('002_S_0413', 0, 'T1'): 1.0,
                  ('002_S_4213', 2, 'DTI'): 1.0,
                  ('002_S_4213', 2, 'T1'): 1.0,
                  ('002_S_4799', 0, 'DTI'): 1.0,
                  ('002_S_4799', 0, 'T1'): 1.0,
                  ('002_S_5178', 0, 'DTI'): 1.0,
                  ('002_S_5178', 0, 'T1'): 1.0,
                  ('002_S_5230', 2, 'DTI'): 1.0,
                  ('002_S_5230', 2, 'T1'): 1.0}
})
df1.index.names = ['Subject ID', 'Visit_NUM', 'Description']

# df2：普通结构，记录每个Subject的多次Visit_NUM
df2 = pd.DataFrame({
    'Subject ID': {0: '002_S_0413',
                   1: '002_S_0413',
                   2: '002_S_4213',
                   3: '002_S_4213',
                   4: '002_S_4799',
                   5: '002_S_4799',
                   6: '002_S_4799',
                   7: '002_S_5178',
                   8: '002_S_5178',
                   9: '002_S_5230',
                   10: '002_S_5230',
                   11: '002_S_5230',
                   12: '002_S_6007',
                   13: '002_S_6007'},
    'Visit_NUM': {0: 0,
                  1: 2,
                  2: 0,
                  3: 2,
                  4: 0,
                  5: 1,
                  6: 2,
                  7: 0,
                  8: 2,
                  9: 0,
                  10: 1,
                  11: 2,
                  12: 0,
                  13: 1}
})

现有代码（仅支持等于匹配） #

df3 = pd.DataFrame(df1.reset_index([
    'Visit_NUM', 'Description']).groupby(
    level=0)['Visit_NUM'].transform(lambda x: x + 2)).reset_index(
).drop_duplicates(['Subject ID'])

t = df3.merge(df2.reset_index(), on=['Subject ID', 'Visit_NUM'])
t = t['Subject ID']
out = df1.loc[df1.index.get_level_values('Subject ID').isin(t)]

优化后的简洁实现 #

核心思路 #

1. 提取df1中每个Subject ID对应的唯一Visit_NUM（同一Subject的Visit_NUM在df1中一致）
2. 提取df2中每个Subject ID的最大Visit_NUM（若最大值满足≥条件，则必然存在符合要求的记录）
3. 筛选出满足df2最大Visit_NUM ≥ df1 Visit_NUM +2的Subject ID
4. 保留df1中符合条件的所有行

代码实现 #

# 提取df1中每个Subject的唯一Visit_NUM
df1_subj_visit = df1.reset_index()[['Subject ID', 'Visit_NUM']].drop_duplicates()

# 计算df2中每个Subject的最大Visit_NUM
df2_subj_max_visit = df2.groupby('Subject ID')['Visit_NUM'].max().reset_index()

# 合并并筛选符合条件的Subject ID
qualified_subjects = df1_subj_visit.merge(df2_subj_max_visit, on='Subject ID')\
    .query('Visit_NUM_y >= Visit_NUM_x + 2')['Subject ID']

# 筛选df1结果
out = df1[df1.index.get_level_values('Subject ID').isin(qualified_subjects)]

期望输出 #

print(out.to_dict())
# 输出结果：
{'Sex': {('002_S_0413', 0, 'DTI'): 'F',
         ('002_S_0413', 0, 'T1'): 'F',
         ('002_S_4799', 0, 'DTI'): 'M',
         ('002_S_4799', 0, 'T1'): 'M',
         ('002_S_5178', 0, 'DTI'): 'M',
         ('002_S_5178', 0, 'T1'): 'M'},
 'DIAGNOSIS': {('002_S_0413', 0, 'DTI'): 1.0,
               ('002_S_0413', 0, 'T1'): 1.0,
               ('002_S_4799', 0, 'DTI'): 1.0,
               ('002_S_4799', 0, 'T1'): 1.0,
               ('002_S_5178', 0, 'DTI'): 1.0,
               ('002_S_5178', 0, 'T1'): 1.0}}

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