创建掩码以识别 Pandas 数据框中组中的最后两行

Question

我有一个包含分组变量的 Pandas 数据框。可以使用以下方法生成一个示例：

df = pd.DataFrame({'grp':['a','a','b','b','b','c','d','d','d','d'],
                   'data':[4,5,3,6,7,8,9,8,7,3]})

...看起来像：

  grp  data
0   a     4
1   a     5
2   b     3
3   b     6
4   b     7
5   c     8
6   d     9
7   d     8
8   d     7
9   d     3

我可以使用以下方法检索每组的最后两行：

dfgrp = df.groupby('grp').tail(2)

但是，我想生成一个掩码来标识最后两行（如果仅存在 1 行，则为 1 行），理想情况下会生成如下所示的输出：

0     True
1     True
2    False
3     True
4     True
5     True
6    False
7    False
8     True
9     True

我认为这相对简单，但我一直无法找到解决方案。如有建议，我们将不胜感激。

Answer 1

如果您的索引是唯一的，您可以使用 来做到这一点isin。

import pandas as pd

df = pd.DataFrame({'grp':['a','a','b','b','b','c','d','d','d','d'],
                   'data':[4,5,3,6,7,8,9,8,7,3]})
df['mask'] = df.index.isin(df.groupby('grp').tail(2).index)
df

grp data    mask
0   a   4   True
1   a   5   True
2   b   3   False
3   b   6   True
4   b   7   True
5   c   8   True
6   d   9   False
7   d   8   False
8   d   7   True
9   d   3   True