Sum of rows based on index with Numpy

Question

I have a 2D array composed of 2D vectors and a 1D array of indices.

How can I add / sumvthe rows of the 2D array that share the same index, using numpy?

Example:

arr = np.array([[48, -51], [-15, -55], [26, -49], [-13, -17], [-67,  -7], [23, -48], [-29, -64], [37,  68]])
idx = np.array([0, 1, 1, 2, 2, 3, 3, 4])

#desired output
array([[48, -51],
       [11, -104],
       [-80, -24],
       [-6, -112],
       [ 37, 68]])

Notice how the original array arr is of shape (8, 2), and the result of the operation is (5, 2).

Not sure about the "always", but most of the time certainly. — solub
– solub, Commented Oct 24, 2019 at 19:02

Mad Physicist · Accepted Answer · 2019-10-24 19:57:18Z

2

If the indices are not always grouped, apply np.argsort first:

order = np.argsort(idx)

You can compute the locations of the sums using np.diff followed by np.flatnonzero to get the indices. We'll also prepend zero and shift everything by 1:

breaks = np.flatnonzero(np.concatenate(([1], np.diff(idx[order])))

breaks can now be used as an argument to np.add.reduceat:

result = np.add.reduceat(arr[order, :], breaks, axis=0)

If the indices are already grouped, you don't need to use order at all:

breaks = np.flatnonzero(np.concatenate(([1], np.diff(idx)))
result = np.add.reduceat(arr, breaks, axis=0)

edited Oct 24, 2019 at 19:57

answered Oct 24, 2019 at 19:49

Mad Physicist

116k29 gold badges202 silver badges292 bronze badges

Sign up to request clarification or add additional context in comments.

Comments

Quang Hoang · Accepted Answer · 2019-10-24 19:31:30Z

1

You can use pandas for the purpose:

pd.DataFrame(arr).groupby(idx).sum().to_numpy()

Output:

array([[  48,  -51],
       [  11, -104],
       [ -80,  -24],
       [  -6, -112],
       [  37,   68]])

answered Oct 24, 2019 at 19:31

Quang Hoang

151k11 gold badges64 silver badges86 bronze badges

Collectives™ on Stack Overflow

Sum of rows based on index with Numpy

2 Answers 2

Comments

Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

2 Answers 2

Comments

Comments

Your Answer

Sign up or log in

Post as a guest

Related