在 pandas 中使用 groupby 进行滚动窗口选择

Question

我有以下熊猫数据框：

# Create the DataFrame
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]
})
df

id  date    value
0   1   1   11
1   1   2   12
2   1   3   13
3   1   4   14
4   1   5   15
5   1   6   16
6   1   7   17
7   1   8   18
8   2   5   21
9   2   6   22
10  2   7   23
11  2   8   24
12  2   9   25
13  2   10  26
14  2   11  27
15  2   12  28

我想以滚动窗口的方式查询上面的数据框，以获得两个

ids

。滚动窗口的大小应为

n

。

所以，如果

n==2

，在第一次迭代中我想查询这个：

df.query('(id==1 and (date==1 or date==2)) or (id==2 and (date==5 or date==6))')

id  date    value
0   1   1   11
1   1   2   12
8   2   5   21
9   2   6   22

在第二次迭代中我想查询这个：

df.query('(id==1 and (date==2 or date==3)) or (id==2 and (date==6 or date==7))')

id  date    value
1   1   2   12
2   1   3   13
9   2   6   22
10  2   7   23

在第三次迭代中我想查询这个：

df.query('(id==1 and (date==3 or date==4)) or (id==2 and (date==7 or date==8))')

id  date    value
2   1   3   13
3   1   4   14
10  2   7   23
11  2   8   24

等等。我怎么能在 pandas 中做到这一点？我的数据大约有 500

ids

Answer 1

确切的预期逻辑尚不完全清楚，但假设您想循环遍历组/卷，您可以将

groupby.nth

与

sliding_window_view

结合起来：

from numpy.lib.stride_tricks import sliding_window_view as swv

n = 2

max_size = df['id'].value_counts().max()
g = df.sort_values(by=['id', 'date']).groupby('id')

for idx in swv(np.arange(max_size), 2):
    print(f'rows {idx}')
    print(g.nth(idx))

输出：

rows [0 1]
   id  date  value
0   1     1     11
1   1     2     12
8   2     5     21
9   2     6     22
rows [1 2]
    id  date  value
1    1     2     12
2    1     3     13
9    2     6     22
10   2     7     23
rows [2 3]
    id  date  value
2    1     3     13
3    1     4     14
10   2     7     23
11   2     8     24
rows [3 4]
    id  date  value
3    1     4     14
4    1     5     15
11   2     8     24
12   2     9     25
rows [4 5]
    id  date  value
4    1     5     15
5    1     6     16
12   2     9     25
13   2    10     26
rows [5 6]
    id  date  value
5    1     6     16
6    1     7     17
13   2    10     26
14   2    11     27
rows [6 7]
    id  date  value
6    1     7     17
7    1     8     18
14   2    11     27
15   2    12     28