3

I am looking for a vectorization method for summing all elements of a numpy array that is evenly divisible by 5.

For example, if I have

test = np.array([1,5,12,15,20,22])

I want to return 40. I know about the np.sum method, but is there a way to do this using vectorization given the X%5 == 0 conditional?

SeeDerekEngineer
  • 770
  • 2
  • 6
  • 22

1 Answers1

6

We can use the mask of the matches for boolean-indexing to select those elements and then simply sum those, like so -

test[test%5==0].sum()

Sample step-by-step run -

# Input array
In [48]: test
Out[48]: array([ 1,  5, 12, 15, 20, 22])

# Mask of matches
In [49]: test%5==0
Out[49]: array([False,  True, False,  True,  True, False], dtype=bool)

# Select matching elements off input
In [50]: test[test%5==0]
Out[50]: array([ 5, 15, 20])

# Finally sum those elements
In [51]: test[test%5==0].sum()
Out[51]: 40
Divakar
  • 218,885
  • 19
  • 262
  • 358