python使用滚动函数rolling()连续选择数据

对于二维数组，有时候需要类似于滑动窗口连续滚动选择数据，本文使用rolling()函数滚动选择数据。

# -*- coding: utf-8 -*-
import pandas as pd
import numpy as np

# 设置一个初始数据框
df1 = [1,2,3,4,5]
df2 = [2,3,4,5,6]
df3 = [3,4,5,6,7]
df4 = [4,5,6,7,8]
df = pd.DataFrame({'a':list(df1),'b':list(df2),'c':list(df3),'d':list(df4)})
print("input data:\n",df)

# 多列滚动函数
# handle对滚动的数据框进行处理

df_result = []
def handle(x,df,n):
  if x+n < df.shape[1] or x+n == df.shape[1]:      #当滚动框滚动到数据边界位置停止滚动
    df = df[:].iloc[ : , x : x+n]                  #使用iloc函数选择需要滚动操作的数据        
    df = np.array(df)
    df_result.append(df)                           #保存滚动操作选择的数据
    print("roll result:\n",df)
  return 1

# group_rolling 进行滚动 
# n：滚动的行数
# df：目标数据框
# name：要滚动的列名
def group_rolling(n,df):
  df_roll = pd.DataFrame({0:list(range(len(df)-n+1))})
  df_roll[0].rolling(window=1).apply(lambda x:handle(int(x[0]),df,n),raw=True)  #window：表示时间窗的大小

group_rolling(n=2,df=df)

df_result1 = np.array(df_result)
print("roll result saved:\n",df_result1)
print("df_result1:",df_result1.shape)

滚动选择的结果如下，将滚动结果保存在df_result1中构成尺寸为（3,5,2）数组。

 input data:
   a  b  c  d
0  1  2  3  4
1  2  3  4  5
2  3  4  5  6
3  4  5  6  7
4  5  6  7  8
roll result:
 [[1 2]
  [2 3]
  [3 4]
  [4 5]
  [5 6]]
roll result:
 [[2 3]
  [3 4]
  [4 5]
  [5 6]
  [6 7]]
roll result:
 [[3 4]
  [4 5]
  [5 6]
  [6 7]
  [7 8]]
roll result saved:
 [[[1 2]
   [2 3]
   [3 4]
   [4 5]
   [5 6]]

 [[2 3]
  [3 4]
  [4 5]
  [5 6]
  [6 7]]

 [[3 4]
  [4 5]
  [5 6]
  [6 7]
  [7 8]]]
df_result1: (3, 5, 2)