如何用NumPy高效按双键分组ID为二维列表数组

现有四个NumPy数组class_label、student_name、student_id和grade，已通过以下代码按班级+学生姓名的组合聚合成绩：

import numpy as np
import numpy_groupies as npg

class_label  = [0,  2, 1, 1, 2, 3, 0, 1]
student_name = [0,  0, 1, 2, 1, 1, 0, 1]
student_id   = [0,  1, 2, 3, 4, 5, 6, 7]
grade        = [10, 5, 7, 4, 9, 8, 9, 3]

n_classes = 4
n_names   = 3

group_idx = np.vstack((class_label, student_name))
# 按班级+姓名聚合成绩的结果
score_agg = npg.aggregate(group_idx, grade, size=(4,3))

上述代码的聚合结果为：

array([[19,  0,  0],
       [ 0, 10,  4],
       [ 5,  9,  0],
       [ 0,  8,  0]])

现在需要在尽量减少代码开销的前提下，构建一个形状为(n_classes, n_names)的二维对象数组，其中每个单元格存储对应(class_label, student_name)分组下的所有student_id组成的列表。期望结果如下：

array([[[0, 6], [], []],
       [[], [2, 7], [3]],
       [[1], [4], []],
       [[], [5], []]], dtype=object)

比如数组第0行第0列（班级0、姓名0）的位置，对应所有class_label==0且student_name==0的记录ID列表[0,6]。

直接复用现有分组逻辑，借助numpy_groupies的aggregate函数，指定聚合规则为收集列表、空分组填充空列表，即可高效生成目标数组：

# 构建分组对应的学生ID对象数组
id_agg = npg.aggregate(group_idx, student_id, size=(n_classes, n_names), func=list, fill_value=[])

关键参数说明 #

• func=list：将每个分组内的student_id元素直接收集为列表
• fill_value=[]：为无匹配记录的分组设置默认空列表，保证数组形状严格符合(4,3)要求

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