Python 如何更改groupby作用域以找到满足掩码条件的第一个值

这是post的扩展.

我的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']
    }
)

这是预期输出.我想创建列result:

   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)
)

过程是这样开始的:

a) groupby列为sub

b)找到满足每个组掩码条件的第一行.

c)在结果中输入label的值

如果没有满足掩码条件的行，则groupby列更改为main以查找掩码的第一行.这一阶段有条件:

当使用main作为groupby列时，不应考虑之前的subs.

子列中的组d的上述步骤示例:

a) sub是groupby列.

b)在d组中没有df.num_2在df.num_1和df.num_3之间的行(mask的条件)

所以现在对于组d，搜索其主组.第c章也是这个主要群体.由于它在组d之前，组c不应计入该步骤.因此，在x组中，mask的第一行有f标签(c d 222).<<

需要注意的是，对于每sub个组，num_2在整个组中不会改变.例如，对于整个组c，num_2是sub.

这是我基于这answer的try ，但它不工作:

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)