numpy: efficiently summing with index arrays

For completeness, in numpy >= 1.8 you can also use np.add‘s at method:

In [8]: m, n = np.random.rand(2, 10)

In [9]: m_idx, n_idx = np.random.randint(10, size=(2, 20))

In [10]: m0 = m.copy()

In [11]: np.add.at(m, m_idx, n[n_idx])

In [13]: m0 += np.bincount(m_idx, weights=n[n_idx], minlength=len(m))

In [14]: np.allclose(m, m0)
Out[14]: True

In [15]: %timeit np.add.at(m, m_idx, n[n_idx])
100000 loops, best of 3: 9.49 us per loop

In [16]: %timeit np.bincount(m_idx, weights=n[n_idx], minlength=len(m))
1000000 loops, best of 3: 1.54 us per loop

Aside of the obvious performance disadvantage, it has a couple of advantages:

  1. np.bincount converts its weights to double precision floats, .at will operate with you array’s native type. This makes it the simplest option for dealing e.g. with complex numbers.
  2. np.bincount only adds weights together, you have an at method for all ufuncs, so you can repeatedly multiply, or logical_and, or whatever you feel like.

But for your use case, np.bincount is probably the way to go.

Leave a Comment