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

Question

Допустим, у меня есть этот кадр данных для панд:

import pandas as pd
import numpy as np
df = pd.DataFrame({'x': np.random.randint(-10, 10, size=100),
                   'y': np.random.randint(-10, 10, size=100)})

И у меня есть любой запрос, который выбирает строки, например,

query = (df['x'] > 3) & (df['y'] < 0)

Как получить группы строк, которые соответствуют этомузапрос И следующие последовательные k строк (если их меньше k, возвращать, сколько их доступно)?

Например, для k = 2 громоздкий и ручной способ сделать это:

# 1st value
sel0 = df[query].reset_index()

# 2nd value
sel1 = df[query.shift(1).fillna(False)].reset_index()

# 3rd value
sel2 = df[query.shift(2).fillna(False)].reset_index()

concat_df = pd.concat([sel0, sel1, sel2])
grouped_df = concat_df.groupby(concat_df.index)
groups = [grouped_df.get_group(i) for i in grouped_df.groups]

Существует ли однострочник, который может обобщить это для любого k и выполнить его быстро?

Answer 1

Я думаю, что вы можете сделать это, используя cumsum, groupby и head:

Попробуйте, где k = 2 использует head (3), текущая запись плюс два:

df.groupby(query.cumsum()).head(3)

и обобщить попробуйте это

k=2
df.groupby(query.cumsum()).head(k+1)