python-3.x - 如何从 Pandas 数据框中的组中获取行的子集？

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