如何快速检索numpy数组并对指定ID对应的数值求和

性能问题原因 #

你当前的实现每次查询都要遍历全量16万条数据做ID匹配，3000次查询的时间复杂度为O(总数据量×查询ID数)，是耗时过高的核心原因。改为一次性完成全量数据分组求和、再按需查询结果的逻辑，速度会提升多个量级。

优化方案 #

方案1：纯NumPy实现（无额外依赖）

仅需一次全量预处理，后续3000次查询均为O(1)时间复杂度，总耗时可降至1秒以内：

# 预处理步骤 仅执行一次
ids = data_arr[:, 0]
# 提前把数值列转为浮点型，避免运算时的类型转换开销
values = data_arr[:, 1].astype(np.float64)
# 获取所有唯一ID和对应分组编码
unique_ids, group_idx = np.unique(ids, return_inverse=True)
# 按分组批量求和
sum_per_id = np.bincount(group_idx, weights=values)
# 转成字典方便查询
sum_dict = dict(zip(unique_ids, sum_per_id))

# 求和函数直接查字典即可
def sumValues(self, id):
    return sum_dict.get(id, 0.0)

方案2：Pandas实现（代码更简洁）

如果允许引入Pandas依赖，分组逻辑写法更简单，也能达到相近的性能：

import pandas as pd

# 预处理步骤 仅执行一次
df = pd.DataFrame(data_arr, columns=["id", "value"])
df["value"] = df["value"].astype(float)
sum_dict = df.groupby("id")["value"].sum().to_dict()

# 求和函数和方案1一致
def sumValues(self, id):
    return sum_dict.get(id, 0.0)

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