python - Pandas:查找与组中谓词匹配的第一行的组索引(如果有)

Question

我想按某些条件对 DataFrame 进行分组，然后在组中找到满足某些谓词的第一行的整数索引(不是 DataFrame)。如果没有这样的行，我想得到 NaN。

例如，我按列 a 除以 5 进行分组，然后在每个组中查找列 b 为“红色”的第一行的索引:

import pandas as pd
import numpy as np

df = pd.DataFrame({'a': xrange(12), 'b': ['red', 'green', 'blue'] * 4})


     a      b
0    0    red
1    1  green
2    2   blue
3    3    red
4    4  green
5    5   blue
6    6    red
7    7  green
8    8   blue
9    9    red
10  10  green
11  11   blue

df.groupby(df.a // 5).apply(lambda g: next((idx for idx, row in g.reset_index(drop=True).iterrows() if row.b == "red"), None))


a
0     0
1     1
2   NaN
dtype: float64

(我想我假设行的顺序与原始 DataFrame 中的顺序相同，但如果需要，我可以对组进行排序。)是否有更简洁、有效的方法来执行此操作？

Answer 1

这有点长，但恕我直言，更容易理解/可定制

In [126]: df2 = df.copy()

这是您的群组指标

In [127]: g = df.a//5

对创建组的引用

In [128]: grp = df.groupby(g)

创建生成的组和组内累积计数的列

In [129]: df2['group'] = g

In [130]: df2['count'] = grp.cumcount()

In [131]: df2
Out[131]: 
     a      b  group  count
0    0    red      0      0
1    1  green      0      1
2    2   blue      0      2
3    3    red      0      3
4    4  green      0      4
5    5   blue      1      0
6    6    red      1      1
7    7  green      1      2
8    8   blue      1      3
9    9    red      1      4
10  10  green      2      0
11  11   blue      2      1

过滤和分组可以返回您想要的第一个元素。计数是组内计数

In [132]: df2[df2.b=='red'].groupby('group').first()
Out[132]: 
       a    b  count
group               
0      0  red      0
1      6  red      1

您可以生成所有组 key (例如，您的过滤器没有返回任何内容)；这边走。

In [133]: df2[df2.b=='red'].groupby('group').first().reindex(grp.groups.keys())
Out[133]: 
    a    b  count
0   0  red      0
1   6  red      1
2 NaN  NaN    NaN