Как получить подмножество строк из группы в кадре данных панд?

Question

У меня есть фрейм данных с идентификатором столбца и двоичным столбцом, как в примере ниже

     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

Что я хочу сделать, так это получить первые четыре последовательных записи, которые равны 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

Индекс внутри группы, где есть четыре последовательных, отличается для каждой группы, как в примере. Как мне это сделать?

Я попробовал решение, предложенное Биллом Дж. В этом вопросе , но у меня это не сработало.

Работа с фреймами данных Pandas и Python 3.6

Answer 1

query и groupby с 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 ed значений по сравнению 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 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