如何更改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': [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) 用sub作为分组的列。
b) 找到每个组中符合掩码条件的第一行。
c) 将label的值放入结果中。
如果没有行符合掩码的条件，那么就把分组列改成main，去找掩码的第一行。在这个阶段有个条件：
之前的sub组在用main作为分组列时不应该被考虑。
下面是对上述步骤的一个示例，针对d组在sub列中的情况：
a) sub是分组列。
b) 在d组中，没有行的df.num_2在df.num_1和df.num_3之间（这是掩码的条件）。
所以现在要为d组搜索它的主组。不过c组也在这个主组中。由于c组在d组之前，所以c组在这一步不算。因此在x组中，掩码的第一行有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)
1 个回答
1
你可以把我之前的代码更新一下，使用两个比较。要注意的是，要把 num_2 当作列来使用。此外，你还需要把输出的参考列改成“label”：
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