如何调整groupby范围，按规则获取符合掩码条件的首个label值？

数据与需求 #

原始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': [97, 90, 105, 2100, 1000, 101, 110, 222, 90, 100, 99, 90, 2, 92, 95, 93],
        'num_2': [100, 100, 100, 102, 102, 209, 209, 209, 209, 100, 100, 100, 100, 90, 90, 90],
        'num_3': [99, 110, 110, 110, 110, 222, 222, 222, 222, 150, 101, 200, 5, 95, 95, 100],
        'label': ['a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k', 'l', 'm', 'n', 'o', 'p']
    }
)

预期输出 #

main sub  num_1  num_2  num_3 label result
0     x   c     97    100     99     a      b
1     x   c     90    100    110     b      b
2     x   c    105    100    110     c      b
3     x   d   2100    102    110     d      f
4     x   d   1000    102    110     e      f
5     x   e    101    209    222     f      f
6     x   e    110    209    222     g      f
7     x   e    222    209    222     h      f
8     x   e     90    209    222     i      f
9     y   f    100    100    150     j      k
10    y   f     99    100    101     k      k
11    y   f     90    100    200     l      k
12    y   f      2    100      5     m      k
13    y   g     92     90     95     n    NaN
14    y   g     95     90     95     o    NaN
15    y   g     93     90    100     p    NaN

掩码条件 #

mask = (
    (df.num_1 < df.num_2) & 
    (df.num_2 < df.num_3)
)

核心逻辑 #

1. 优先按sub分组：在每个sub组内，找到第一个满足掩码条件的行，将其label填充到该组所有行的result列
2. 兜底按main分组：若当前sub组无符合条件的行，则在对应的main组中查找，但需排除当前sub组之前的所有sub组；找到首个满足掩码条件的行后填充其label，若无则填充NaN

示例：sub组d无符合条件的行，切换到main组x后排除sub组c，找到sub组e中首个满足条件的行（label为f），填充到sub组d的所有行。

错误尝试代码 #

def find(g):
    # get sub as 0,1,2…
    sub = pd.factorize(g['sub'])[0]
    # convert inputs to numpy
    a = g['num_1'].to_numpy()
    b = g.loc[~g['sub'].duplicated(), 'num_2'].to_numpy()
    c = g['num_3'].to_numpy()
    # form mask
    # (a[:, None] > b) -> num_1 > num_2
    # (sub[:, None] >= np.arange(len(b))) -> exclude previous groups
    m = (a[:, None] < b) &  (a[:, None] > c) & (sub[:, None] >= np.arange(len(b)))
    # find first True per column
    return pd.Series(np.where(m.any(0), a[m.argmax(0)], np.nan)[sub],
                     index=g.index)

df['result'] = df.groupby('main', group_keys=False).apply(find)

正确实现方法 #

import numpy as np

def process_main_group(main_group):
    # 给每个sub组分配唯一序号，按出现顺序排序
    sub_labels, sub_indices = pd.factorize(main_group['sub'], sort=False)
    unique_subs = main_group['sub'].unique()
    sub_to_idx = {sub: i for i, sub in enumerate(unique_subs)}
    
    # 预计算所有行的掩码
    mask = (main_group['num_1'] < main_group['num_2']) & (main_group['num_2'] < main_group['num_3'])
    # 保存每个sub组的首个符合条件的label
    sub_result = {}
    
    # 先处理每个sub组自身的情况
    for sub in unique_subs:
        sub_group = main_group[main_group['sub'] == sub]
        # 找当前sub组中第一个满足掩码的行
        first_match = sub_group[mask.loc[sub_group.index]].head(1)
        if not first_match.empty:
            sub_result[sub] = first_match['label'].iloc[0]
    
    # 处理自身无匹配的sub组，在main组后续sub中查找
    for i, sub in enumerate(unique_subs):
        if sub in sub_result:
            continue
        # 排除当前sub之前的所有sub组，只保留后续的sub
        allowed_subs = unique_subs[i+1:]
        if not allowed_subs:
            sub_result[sub] = np.nan
            continue
        # 在后续sub组中找第一个满足掩码的行
        filtered = main_group[main_group['sub'].isin(allowed_subs)]
        first_match = filtered[mask.loc[filtered.index]].head(1)
        sub_result[sub] = first_match['label'].iloc[0] if not first_match.empty else np.nan
    
    # 将结果映射回原main组的每一行
    return main_group['sub'].map(sub_result)

# 按main分组处理，生成result列
df['result'] = df.groupby('main', group_keys=False).apply(process_main_group)

代码说明 #

1. 按main分组处理：每个main组独立计算，避免跨main组干扰
2. sub组序号标记：通过factorize记录sub组的出现顺序，确保排除之前的sub组
3. 两步查找逻辑：
   • 第一步：遍历每个sub组，先找自身是否有符合条件的首行，记录label
   • 第二步：对无匹配的sub组，在当前main组的后续sub组中查找首个符合条件的行
4. 结果映射：将每个sub组对应的结果映射到原DataFrame的每一行

运行上述代码后，即可得到预期的result列。

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