如何结合numpy.unique()高效生成总分？大规模数据性能优化求助

方案1：使用np.bincount（最简洁高效，适合非负整数ID） #

如果你的学生ID都是非负整数，np.bincount是最优选择——它是numpy专门为这类分组求和场景设计的底层优化函数，完全替代循环：

import numpy as np

labels = np.array([
    [0,1,5],
    [0,1,3],
    [2,4,5]])
weights = np.array([
    [1],
    [2],
    [4]])

# 展平标签数组（替代concatenate，更简洁）
c_labels = labels.ravel()
# 广播权重到与标签匹配的形状并展平
c_weights = np.broadcast_to(weights, labels.shape).ravel()

# 直接计算每个学生的加权总分：索引对应学生ID，值为总分
scores = np.bincount(c_labels, weights=c_weights)
# 若需要唯一学生ID数组（与原代码的`uni`一致）
uni = np.unique(c_labels)

print(uni)
print(scores)

输出结果和原代码完全一致：

[0 1 2 3 4 5]
[3. 3. 4. 2. 4. 5.]

方案2：使用np.add.at（通用场景，支持任意ID类型） #

如果学生ID不是连续非负整数（比如存在大跨度ID、字符串ID），可以结合np.unique的索引结果，用np.add.at实现原地向量化累加：

import numpy as np

labels = np.array([
    [0,1,5],
    [0,1,3],
    [2,4,5]])
weights = np.array([
    [1],
    [2],
    [4]])

c_labels = labels.ravel()
c_weights = np.broadcast_to(weights, labels.shape).ravel()

# 获取唯一ID及其对应的索引映射
uni, inv = np.unique(c_labels, return_inverse=True)
# 初始化分数数组，用np.add.at完成批量累加
scores = np.zeros(len(uni))
np.add.at(scores, inv, c_weights)

print(uni)
print(scores)

性能提升原理 #

原代码的for循环是Python层面的逐次操作，每次循环都要执行布尔索引、求和等步骤，大数据量下会产生巨大的性能开销。而np.bincount和np.add.at都是C层面实现的向量化操作，直接对整个数组进行批量计算，效率能提升几个数量级，完全适配大规模数据的重复执行场景。

内容的提问来源于stack exchange，提问作者Kay K.