Python | NumPy 相同x值对应y值求和算法性能低下问题求解

性能瓶颈分析 #

原实现的时间复杂度为O(N*K)（N为总数据行数，K为唯一x的数量），每次循环都对全量数据做布尔匹配，Python级循环+重复全表扫描是速度慢的核心原因。
另外原实现存在逻辑偏差：data[data[:,0] == item].sum()会把x列的值也一同累加，若需求仅为累加对应y值，需调整为对第二列单独求和，以下方案均按仅累加y值的需求实现，可根据需要调整。

方案1：pandas groupby实现（代码简洁易维护） #

万级数据处理耗时在毫秒级，代码可读性极高，无需手动实现分组逻辑：

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

data = np.array([[1, 2], [1, 3], [1, 5], [2, 2], [2, 4], [3, 1], [3, 8], [3, 9]])
# 按x列分组，对y列求和
group_res = pd.DataFrame(data, columns=['x', 'y']).groupby('x')['y'].sum()
x = group_res.index.values
y = group_res.values

plt.plot(x, y)
plt.show()

方案2：纯numpy原生实现（无额外依赖，性能最高） #

全程无Python级循环，所有运算均在C层面执行，性能比原实现提升百倍以上：

import numpy as np
import matplotlib.pyplot as plt

data = np.array([[1, 2], [1, 3], [1, 5], [2, 2], [2, 4], [3, 1], [3, 8], [3, 9]])
# return_inverse返回原数组每个x对应的唯一值索引
x_unique, inverse_idx = np.unique(data[:, 0], return_inverse=True)
# 按索引对y值加权求和
y_sum = np.bincount(inverse_idx, weights=data[:, 1])

plt.plot(x_unique, y_sum)
plt.show()

如果需要和原实现逻辑保持一致（x和y值一同累加），只需要将weights=data[:,1]修改为weights=data[:,0] + data[:,1]即可。

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