python - NumPy 根据另一个数组中的值对第三个数组中的每个匹配元素求和一个数组

Question

我有两个 numpy 数组，一个包含值，一个包含每个值类别。

values=np.array([1,2,3,4,5,6,7,8,9,10])
valcats=np.array([101,301,201,201,102,302,302,202,102,301])

我有另一个数组，其中包含我想要汇总的独特类别。

categories=np.array([101,102,201,202,301,302])

我的问题是，我将运行相同的求和过程数十亿次，每一微秒都很重要。

我目前的实现如下。

catsums=[]
for x in categories:
    catsums.append(np.sum(values[np.where(valcats==x)]))

生成的 catsums 应该是:

[1, 14, 7, 8, 12, 13]

我当前的运行时间约为 5 微秒。我对 Python 还是有些陌生，希望通过潜在地组合前两个数组或 lamdba 或一些我什至不知道的很酷的东西来找到一个快速的解决方案。

感谢阅读!

Answer 1

您可以使用searchsorted 和bincount -

np.bincount(np.searchsorted(categories, valcats), values)