如何更改 groupby 范围以查找满足掩码条件的第一个值？

Question

这是此帖子的扩展。

我的数据框是：

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

的值放入结果中

如果没有满足mask条件的行，则将groupby列更改为

main

以查找mask的第一行。此阶段有条件：

使用

sub

作为

main

列时，不应考虑之前的

groupby

。

子栏中

d

组的上述步骤示例：

a)

sub

是 groupby 列。

b)

d

组中没有

df.num_2

位于

df.num_1

和

df.num_3

之间的行（

mask

的条件）

所以现在对于组

d

，搜索其主要组。然而，

c

组也属于这个主要组。由于它位于组

d

之前，因此组

c

不应计入此步骤。因此，在

x

组中，

mask

的第一行具有

f

标签 (101 < 102 < 222).

需要注意的一件事是，对于每个

sub

组，

num_2

在整个组中不会发生变化。例如，对于整个组

c

num_2

是 100。

这是我基于此答案的尝试，但它不起作用：

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)

Answer 1

您可以更新我之前的代码以使用两次比较，您必须注意使用 num_2 作为列。您还需要将输出参考列更改为“标签”：

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

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

输出：

   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