CODEWITHSUNDEEP
numpyvectorization
Bram
Bram

Reputation: 628

How do I vectorize this calculation in numpy

How do I vectorize the following calculation if a and b are appropriately sized numpy arrays?

total = a[0]
for ix in range(1, len(a)):
   total = total*b[ix-1] + a[ix]

Upvotes: 3

Views: 83

Answers (1)

Daniel F
Daniel F

Reputation: 14399

Never mind, there is a ufunc trick that works if you do some algebra. In this case the ufunc is multiply and the trick is accumulate.

c = np.r_[np.multiply.accumulate(b[0:-1][::-1])[::-1], 1]
total2 = np.sum(a * c)

What this does: Algebraically, you are summing a[i] times the product of b[i:] over for i in range(a.size). To do this, flip b and get the running product of all but the last number (assuming a and b are the same length) and then flip it back. The last value should be 1 because the final a value is just added.

Testing

a = np.random.randint(1, 10, 40)
b = 1 + np.random.rand(40)

total = a[0]
for ix in range(1, len(a)):
   total = total*b[ix-1] + a[ix]

total
278443993.10494208

total2
278443993.10494208

Upvotes: 4

Related Questions