如何高效将DataFrame固定行数的值转为指定形状的列表？

原始数据 #

首先是我们的初始DataFrame：

import pandas as pd
import numpy as np

df5 = pd.DataFrame({'timestamp':['2022-10-01 01:00:00', '2022-10-01 02:00:00',
                                '2022-10-01 03:00:00', '2022-10-01 04:00:00',
                                '2022-10-02 01:00:00', '2022-10-02 02:00:00',
                                '2022-10-02 03:00:00', '2022-10-02 04:00:00'],
                   'A': [1,2,3,4,5,6,7,8],
                   'B': [10,9,8,7,6,5,4,3]}
                  )
df5['timestamp'] = df5['timestamp'].astype('datetime64')

原始数据结构如下：

需求说明 #

我们需要将上述数据按指定行数（示例为4行，可自定义为20行等）分组，把每组内的A、B列值分别打包成列表，生成如下结构的目标DataFrame X（形状为(样本数, 变量数, 时间点数)，示例中是(4,1,4)）：

要求用非迭代的高效方法，避免逐行迭代带来的性能问题。

高效实现方案 #

利用Pandas的向量化操作和分组聚合功能，步骤如下：

1. 生成分组标识：根据自定义的每组行数group_size，给每行数据分配对应的分组ID
2. 转换为长格式：把宽表转为长表，让A、B列的变量名和对应值分别成列
3. 分组聚合：按分组ID和变量名分组，将每组的值聚合为列表
4. 整理结果：调整列名和索引，得到目标结构

完整代码 #

# 自定义每组行数，示例为4，可修改为20等
group_size = 4

# 步骤1：添加分组ID列
df5['group_id'] = np.arange(len(df5)) // group_size

# 步骤2：转长格式
melted_df = df5.melt(id_vars=['group_id'], value_vars=['A', 'B'], var_name='target', value_name='value')

# 步骤3：分组聚合为列表
result = melted_df.groupby(['group_id', 'target'])['value'].agg(list).reset_index(name='values')

# 步骤4：调整索引并移除group_id列（和示例结构完全对齐）
result = result.drop('group_id', axis=1).reset_index(drop=True)

print(result)

运行结果 #

target     values
0      A  [1, 2, 3, 4]
1      B  [10, 9, 8, 7]
2      A  [5, 6, 7, 8]
3      B  [6, 5, 4, 3]

这个方法全程用Pandas内置的向量化操作，没有循环迭代，处理大数据量时效率远高于逐行遍历。

内容的提问来源于stack exchange，提问作者AbelAI