Python numpy array sum over certain indices

Question

How to perform a sum just for a list of indices over numpy array, e.g., if I have an array a = [1,2,3,4] and a list of indices to sum, indices = [0, 2] and I want a fast operation to give me the answer 4 because the value for summing value at index 0 and index 2 in a is 4

Solution 1

You can use sum directly after indexing with indices:

a = np.array([1,2,3,4])
indices = [0, 2] 
a[indices].sum()

Solution 2

The accepted a[indices].sum() approach copies data and creates a new array, which might cause problem if the array is large. np.sum actually has an argument to mask out colums, you can just do

np.sum(a, where=[True, False, True, False])

Which doesn't copy any data.

The mask array can be obtained by:

mask = np.full(4, False)
mask[np.array([0,2])] = True

Solution 3

Try:

>>> a = [1,2,3,4]
>>> indices = [0, 2]
>>> sum(a[i] for i in indices)
4

Faster

If you have a lot of numbers and you want high speed, then you need to use numpy:

>>> import numpy as np
>>> a = np.array([1,2,3,4])
>>> a[indices]
array([1, 3])
>>> np.sum(a[indices])
4