将DataFrame向量化为Numpy数组：多学校数据高效转换方案

需求说明 #

现有如下结构的DataFrame，包含3所学校（ID-1、ID-2、ID-3）的分数数据，每所学校对应4行Score-1、Score-2、Score-3的记录。需要为每所学校生成4×3的数组，最终拼接为3×4×3的Numpy数组。原循环实现在数千所学校的场景下效率偏低，需用向量化方法优化。

DataFrame结构如下：

School  Score-1 Score-2 Score-3
ID-1    1       13      25
ID-1    2       14      26
ID-1    3       15      27
ID-1    4       16      28
ID-2    5       17      29
ID-2    6       18      30
ID-2    7       19      31
ID-2    8       20      32
ID-3    9       21      33
ID-3    10      22      34
ID-3    11      23      35
ID-3    12      24      36

原循环实现代码：

import pandas as pd
import numpy as np

school = ['ID-1', 'ID-1', 'ID-1', 'ID-1', 'ID-2', 'ID-2', 'ID-2', 'ID-2', 'ID-3', 'ID-3', 'ID-3', 'ID-3']
Score_1 = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12]
Score_2 = [13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24]
Score_3 = [25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36]

df = pd.DataFrame({'School':school,
          'Score-1':Score_1,
          'Score-2':Score_2,
          'Score-3':Score_3
            })

schools = list(df['School'].unique())

my_array = []
for school in schools:
    a = np.array(df[df['School']==school].iloc[::,1:])
    my_array.append(a)

my_array = np.array(my_array)

高效向量化实现方法 #

方法1：利用pandas GroupBy + agg 向量化处理 #

通过groupby('School')分组后，直接聚合为数组，再转换为三维Numpy数组。这种方法无需手动循环，完全利用pandas和numpy的内部优化，适合任意每组行数固定的场景。

代码实现：

# 按学校分组，将每组的分数列转为数组，再拼接成三维数组
my_array = df.groupby('School')[['Score-1', 'Score-2', 'Score-3']].agg(lambda x: x.to_numpy()).to_numpy()

方法2：直接用Numpy reshape（仅适用于每组行数固定且数据按学校连续排列） #

如果数据已经按学校分组且每组行数固定（如本例中每所学校4行），可以直接提取分数列的数组后，通过reshape一步到位，这是效率最高的方式。

代码实现：

# 提取分数列的二维数组，reshape为3(学校数) ×4(每组行数) ×3(分数列数)
scores = df[['Score-1', 'Score-2', 'Score-3']].to_numpy()
my_array = scores.reshape(3, 4, 3)

方法对比 #

• 方法1：通用性强，无需依赖数据的排列顺序，适合大多数场景，即使后续每组行数有调整也能兼容。
• 方法2：性能最优，完全基于numpy的底层操作，但要求数据必须按学校连续排列且每组行数严格一致，适合已知数据结构的大规模场景。

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