Pandas中基于GroupBy的滚动窗口数据选取实现方法

嘿，我来给你捋捋怎么高效实现这个需求～ 你的核心需求是：针对每个id的连续日期序列，按固定大小n的滚动窗口，每次迭代把所有id对应位置的窗口行组合起来，对吧？下面给你两种实用的实现方式，尤其适合你有500个id的场景，效率比反复用query高多了！

方法一：分组后直接按索引切片（推荐） #

这个方法思路很直接：先把数据按id分组，然后计算出可以迭代的次数，最后每次遍历都取每个id对应位置的窗口行，合并成结果。优点是代码简洁、执行高效，避免了全表扫描。

完整代码 #

import pandas as pd

# 你的原始数据
df = pd.DataFrame({
    'id': [1, 1, 1, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2, 2, 2],
    'date': [1, 2, 3, 4, 5, 6, 7, 8, 5, 6, 7, 8, 9, 10, 11, 12],
    'value': [11, 12, 13, 14, 15, 16, 17, 18, 21, 22, 23, 24, 25, 26, 27, 28]
})

n = 2  # 滚动窗口大小

# 1. 按id分组，每个组按date排序（保险起见，避免原始数据乱序）
id_groups = {
    group_id: group.sort_values('date').reset_index(drop=True)
    for group_id, group in df.groupby('id')
}

# 2. 计算最大迭代次数：取所有id的窗口数的最小值（确保每个id都有对应窗口）
num_iterations = min(len(group) - n + 1 for group in id_groups.values())

# 3. 遍历每个窗口位置，生成迭代结果
for k in range(num_iterations):
    # 收集当前迭代中每个id的窗口行
    current_window_rows = []
    for group in id_groups.values():
        # 取第k到k+n-1行（对应第k个滚动窗口）
        window_slice = group.iloc[k:k+n]
        current_window_rows.append(window_slice)
    
    # 合并所有id的窗口行
    result_df = pd.concat(current_window_rows, ignore_index=True)
    print(f"第{k+1}次迭代结果：")
    print(result_df)
    print("-" * 30)

代码说明 #

• 分组时先按date排序：如果你的原始数据中每个id的date已经是连续递增的，这一步可以省略，但加上更稳妥；
• 迭代次数计算：len(group) - n + 1是单个id的滚动窗口总数，取所有id的最小值，保证每次迭代每个id都有对应的完整窗口；
• 切片取行：直接用iloc定位窗口行，比构造query语句快得多，尤其是数据量较大时（比如500个id）。

方法二：利用GroupBy + Rolling实现 #

如果你更习惯用pandas的rolling接口，也可以用这种方式，先对每个分组生成滚动窗口，再按迭代次数合并结果：

完整代码 #

import pandas as pd
from collections import defaultdict

df = pd.DataFrame({
    'id': [1, 1, 1, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2, 2, 2],
    'date': [1, 2, 3, 4, 5, 6, 7, 8, 5, 6, 7, 8, 9, 10, 11, 12],
    'value': [11, 12, 13, 14, 15, 16, 17, 18, 21, 22, 23, 24, 25, 26, 27, 28]
})

n = 2

# 1. 对每个id分组，生成滚动窗口并收集有效窗口（长度为n的完整窗口）
window_collector = []
for group_id, group in df.groupby('id'):
    # 对分组应用滚动窗口，窗口大小n
    rolling_obj = group.rolling(n)
    # 遍历每个窗口，只保留完整窗口
    for window_idx, (win_index, win_df) in enumerate(rolling_obj):
        if len(win_df) == n:
            window_collector.append((window_idx, win_df))

# 2. 按窗口索引（迭代次数）分组，合并同一迭代的所有id窗口
iter_result = defaultdict(list)
for idx, win_df in window_collector:
    iter_result[idx].append(win_df)

# 3. 输出每次迭代的结果
for iter_num, dfs in iter_result.items():
    result_df = pd.concat(dfs, ignore_index=True)
    print(f"第{iter_num+1}次迭代结果：")
    print(result_df)
    print("-" * 30)

代码说明 #

• 滚动窗口处理：rolling(n)会生成每个分组的滚动窗口，前n-1个窗口长度不足n，所以用len(win_df) == n过滤掉；
• 按窗口索引合并：用defaultdict把同一迭代次数的所有id窗口行收集起来，再合并成结果DataFrame。

两种方法都能完美实现你要的效果，推荐第一种，因为代码更直观、执行效率更高，适合你的大数据量场景～

备注：内容来源于stack exchange，提问作者quant