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

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

为GroupBy.transform 创建辅助系列，将shifted 的shifted 值与ne (!=) 进行比较，并与另一个条件链接，最后按boolean indexing 过滤：

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 2

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 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