如何切换groupby列，在初始分组未找到符合掩码条件的首行时重新查询？

给定的DataFrame #

import pandas as pd
df = pd.DataFrame(
    {
        'main': ['x', 'x', 'x', 'x', 'x', 'x', 'x', 'x', 'x', 'y', 'y', 'y', 'y', 'y', 'y', 'y'],
        'sub': ['c', 'c', 'c', 'd', 'd', 'e', 'e', 'e', 'e', 'f', 'f', 'f', 'f', 'g', 'g', 'g'],
        'num_1': [10, 9, 80, 80, 99, 101, 110, 222, 90, 1, 7, 10, 2, 10, 95, 10],
        'num_2': [99, 99, 99, 102, 102, 209, 209, 209, 209, 100, 100, 100, 100, 90, 90, 90]
    }
)

预期输出 #

main sub  num_1  num_2  result
0     x   c     10     99     101
1     x   c      9     99     101
2     x   c     80     99     101
3     x   d     80    102     110
4     x   d     99    102     110
5     x   e    101    209     222
6     x   e    110    209     222
7     x   e    222    209     222
8     x   e     90    209     222
9     y   f      1    100     NaN
10    y   f      7    100     NaN
11    y   f     10    100     NaN
12    y   f      2    100     NaN
13    y   g     10     90      95
14    y   g     95     90      95
15    y   g     10     90      95

掩码规则 #

mask = (df.num_1 > df.num_2)

处理流程 #

• a) 以sub列为分组键进行groupby
• b) 查找每个分组中首个符合掩码条件的行
• c) 将该行的num_1值填充到整个分组的result列
• 若当前sub分组无符合条件的行，则切换到以main列为分组键查找，但需要排除当前分组之前的所有sub分组，找到首个符合条件的行后，将其num_1值填充到当前sub分组的result列；如果整个main分组（排除前面sub分组后）也没有符合条件的行，result列填NaN

我的尝试（仅部分解决问题） #

def step_a(g):
    mask = (g.num_1 > g.num_2)

    g.loc[mask.cumsum().eq(1) & mask, 'result'] = g.num_1
    g['result'] = g.result.ffill().bfill()
    return g

a = df.groupby('sub').apply(step_a)

完整解决方案 #

注：预期输出与流程描述存在矛盾（例如sub c和sub d的结果未取后续分组中首个符合掩码的222），以下代码严格遵循流程描述实现，若需完全匹配预期输出需调整规则。

import pandas as pd

df = pd.DataFrame(
    {
        'main': ['x', 'x', 'x', 'x', 'x', 'x', 'x', 'x', 'x', 'y', 'y', 'y', 'y', 'y', 'y', 'y'],
        'sub': ['c', 'c', 'c', 'd', 'd', 'e', 'e', 'e', 'e', 'f', 'f', 'f', 'f', 'g', 'g', 'g'],
        'num_1': [10, 9, 80, 80, 99, 101, 110, 222, 90, 1, 7, 10, 2, 10, 95, 10],
        'num_2': [99, 99, 99, 102, 102, 209, 209, 209, 209, 100, 100, 100, 100, 90, 90, 90]
    }
)

mask = df['num_1'] > df['num_2']

# 1. 预计算每个sub分组中首个符合掩码的num_1值
sub_first_valid = df[mask].groupby('sub')['num_1'].first()
sub_result_map = sub_first_valid.to_dict()

# 2. 为每个main下的sub分组标记顺序，用于排除之前的分组
df['sub_order'] = df.groupby('main')['sub'].transform(lambda x: x.rank(method='dense', ascending=True))

# 3. 定义处理每个sub分组的函数
def process_sub_group(g):
    current_sub = g['sub'].iloc[0]
    current_main = g['main'].iloc[0]
    current_order = g['sub_order'].iloc[0]
    
    # 优先使用当前sub分组的结果
    if current_sub in sub_result_map:
        g['result'] = sub_result_map[current_sub]
        return g
    
    # 若当前分组无结果，在所属main中查找顺序靠后的sub分组
    valid_subs = df[(df['main'] == current_main) & (df['sub_order'] > current_order)]['sub'].unique()
    # 遍历后续sub分组，取第一个有结果的
    for sub in valid_subs:
        if sub in sub_result_map:
            g['result'] = sub_result_map[sub]
            return g
    
    # 若后续无符合条件的分组，填充NaN
    g['result'] = pd.NA
    return g

# 4. 应用函数并整理结果
df_result = df.groupby('sub').apply(process_sub_group).drop(columns='sub_order')
print(df_result)

代码说明 #

1. 预计算sub分组结果：先筛选所有符合掩码的行，按sub分组取首个num_1值，存入字典方便快速查找。
2. 标记sub分组顺序：通过rank方法为每个main下的sub分组分配顺序，确保能准确排除当前分组之前的所有sub分组。
3. 分组处理逻辑：
   • 先检查当前sub分组是否有预计算的结果，有则直接填充。
   • 若无，则遍历当前main下顺序靠后的sub分组，找到第一个有结果的分组，复用其值。
   • 若后续分组均无符合条件的行，填充NaN。

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