Pandas：在分组中查找第一个满足条件的行的组索引（如果有）

Question

我想根据一些标准把一个数据表（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

（我想我假设行的顺序和原始数据表是一样的，不过如果需要的话，我可以对组进行排序。）有没有更简洁、更高效的方法来做到这一点？

Answer 1

我能做到的最好效果是：

import itertools as it
df.groupby(df.a // 5).apply(lambda group: next(it.chain(np.where(group.get_values() == "red")[0], [None])))

唯一真正的区别在于使用了 np.where 来处理数值（所以我觉得这样通常会更快），不过你也可以考虑自己写一个 first_where 函数来用。

Answer 2

这段内容稍微长一点，但我觉得更容易理解和自定义。

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

你可以以这种方式生成所有的组键（例如，你的过滤器没有返回任何结果）。

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