假设我有这个pandas数据帧：

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并快速执行？

解决办法：可以使用cumsum，groupby并且head：

试试这个，其中k = 2使用head（3），当前记录加上两个：

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

并概括试试这个

k=2

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