How to efficiently add or multiply every Nth element of a numpy array?

user_11385378 0
Viewed 159

Suppose I have the following numpy array.

A = [1,2,3,4,5,6]

Question Is there a quick way to multiply or add every nth element in A to yield the following arrays?

B = [3*1, 2*4, 3*5, 4*6]
C = [3+1, 2+4, 3+5, 4+6]

I can accomplish this by rolling A and then adding or multiplying it to the original A.

T = np.roll(A,-2)
B = (A*T)[0:4]
C = (A*T)[0:4]

Just wondering if there is a more pythonic/efficient way to accomplish this? I have also looked at np.add.reduceat and np.multiply.reduceat but they do not seem to allow for skipping values.

numpy arrays python
1 Answers

You can do it in the following way:

A = np.array([1, 2, 3, 4, 5, 6])
k = 2
B = A[:-k]
C = A[k:]
print(B * C)
print(B + C)

Output

[ 3  8 15 24]
[ 4  6  8 10]

Cheers.

user_45960780
Related