python - Pandas : How to get groups of each n rows after row matching query?

Question

假设我有这个 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 并快速执行？

Answer 1

我认为你可以使用cumsum、groupby和head来做到这一点:

试试这个，其中 k=2 使用 head(3)，当前记录加二:

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

并概括地尝试这个

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