Joint accumulation of addition and multiplication

user_14923227 0
Viewed 78

I have an array:

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

and I want to do the following calculation:

out = 0
for e in a:
    out *= 3
    out += e

With out as the output (4582 for the given example), is there a nice way to vectorize this? I think einsum can be used, but I couldn't figure how to write it.

numpy python
2 Answers

Personally I would use reduce:

reduce(lambda x, y: x * 3 + y, a)
user_76942630

One approach:

import numpy as np

a = np.array([1, 2, 3, 1, 3, 4, 2, 4])
powers = np.multiply.accumulate(np.repeat(3, len(a) - 1))
res = np.sum(powers[::-1] * a[:-1]) + a[-1]
print(res)

Output

4582

If you expand the loop, you'll notice that you are multiplying each value of a by a power of 3 and then summing the result.

user_40015920
Related