假设我有一个值数组
values = np.array([0.0, 1.0, 2.0, 3.0, 4.0])
和一个索引数组
indices = np.array([0,1,0,2,2])
有没有一种矢量化的方法来对每个唯一索引的值求和?我的意思是获取此代码片段中的indices
矢量化版本:sums
sums = np.zeros(np.max(indices)+1)
for index, value in zip(indices, values):
sums[index] += value
如果解决方案允许values
(并且因此sums
)是多维的,则可获得加分。
编辑:我对发布的解决方案进行了基准测试:
import numpy as np
import time
import pandas as pd
values = np.arange(1_000_000, dtype=float)
rng = np.random.default_rng(0)
indices = rng.integers(0, 1000, size=1_000_000)
N = 100
now = time.time_ns()
for _ in range(N):
sums = np.bincount(indices, weights=values, minlength=1000)
print(f"np.bincount: {(time.time_ns() - now) * 1e-6 / N:.3f} ms")
now = time.time_ns()
for _ in range(N):
sums = np.zeros(1 + np.amax(indices), dtype=values.dtype)
np.add.at(sums, indices, values)
print(f"np.add.at: {(time.time_ns() - now) * 1e-6 / N:.3f} ms")
now = time.time_ns()
for _ in range(N):
pd.Series(values).groupby(indices).sum().values
print(f"pd.groupby: {(time.time_ns() - now) * 1e-6 / N:.3f} ms")
now = time.time_ns()
for _ in range(N):
sums = np.zeros(np.max(indices)+1)
for index, value in zip(indices, values):
sums[index] += value
print(f"Loop: {(time.time_ns() - now) * 1e-6 / N:.3f} ms")
结果:
np.bincount: 1.129 ms
np.add.at: 0.763 ms
pd.groupby: 5.215 ms
Loop: 196.633 ms
您可以将索引视为箱,将值视为权重
np.bincount
。另一个可能的解决方案是:
首先,创建一个零数组,其长度等于数组
b
中唯一元素的数量indices
然后,它使用该
np.add.at
函数将数组中的值累积到数组指定values
的相应位置。b
indices
输出:
你可以尝试
pd.groupby
由此得出