主页 / user-7973301

Padix Key's questions

Martin Hope
Padix Key
Asked: 2024-10-30 18:25:13 +0800 CST
  • 8

假设我有一个值数组

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
python