如何从熊猫数据框中的组中获取行的子集？

Question

我有一个带有列ID和二进制列的数据框，如下面的示例

     ID    BINARY_MASK
0   101        1
1   101        0
2   101        1
3   101        1
4   101        1
5   101        1
6   101        0
7   101        1
8   102        1 
9   102        1
11  102        1
12  102        1
13  102        0 
14  102        0

我想要做的是获取每个ID组的前四个连续条目，即1。。 我想看到的结果如下：

     ID    BINARY_MASK
2   101        1
3   101        1
4   101        1
5   101        1
8   102        1 
9   102        1
11  102        1
12  102        1

与示例相同，每个组中有四个连续的索引的组内索引不同。 我该怎么做？

我已经尝试了by Bill G in this question提供的解决方案，但是对我来说不起作用。

使用Pandas数据框和Python 3.6

Answer 1

query和groupby with head

最简单的方法是在分组之前过滤掉那些。您可以通过几种方式进行过滤，我选择使用query。

df.query('BINARY_MASK == 1').groupby('ID').head(4)

     ID  BINARY_MASK
0   101            1
2   101            1
3   101            1
4   101            1
8   102            1
9   102            1
11  102            1
12  102            1

Answer 2

为GroupBy.transform创建辅助系列，其中cumsum个值的shift与ne（!=）进行比较，并与另一个条件链接，最后由{ {3}}：

s = df['BINARY_MASK'].ne(df['BINARY_MASK'].shift()).cumsum()
m1 = df.groupby(s)['BINARY_MASK'].transform('size') >= 4
m2 = df['BINARY_MASK'] == 1

df = df[m1 & m2]
print (df)
     ID  BINARY_MASK
2   101            1
3   101            1
4   101            1
5   101            1
7   101            1
8   102            1
9   102            1
11  102            1
12  102            1

Answer 3

使用groupby + head：

df[df['BINARY_MASK']==1].groupby('ID').head(4)

     ID  BINARY_MASK
0   101            1
2   101            1
3   101            1
4   101            1
8   102            1
9   102            1
11  102            1
12  102            1