CODEWITHSUNDEEP
pythonnumpymatrixvectorization
CodeCabbie
CodeCabbie

Reputation: 3454

Sum nth columns elementwise in Numpy matrix

I have the following numpy matrix:

A = [[a,b,c,d,e,f],
     [g,h,i,j,k,l],
     ...]

I need to sum all the nth columns in a element-wise fashion. So, if n is 2, the answer needs to be: 

B = [[a+c+e, b+d+f],
     [g+i+k, h+j+l],
     ...]

(like breaking the matrix into 3, each with 2 columns and adding them.)

But if n is 3, the answer needs to be:

C = [[a+d, b+e, c+f],
     [g+j, h+k, i+l],
     ...]

(like breaking the matrix into 2, each with 3 columns and adding them.)

Is there a general case that will accept the value n without resorting to looping?

Upvotes: 1

Views: 239

Answers (1)

Divakar
Divakar

Reputation: 221604

Reshape to split the last axis into two with the latter of length n and sum along the former -

A.reshape(A.shape[0],-1,n).sum(1)

Sample run -

In [38]: A
Out[38]: 
array([[0, 5, 3, 2, 5, 6],
       [6, 1, 0, 8, 4, 0],
       [8, 6, 1, 5, 7, 0]])

In [39]: n = 2

In [40]: A.reshape(A.shape[0],-1,n).sum(1)
Out[40]: 
array([[ 8, 13],
       [10,  9],
       [16, 11]])

In [41]: n = 3

In [42]: A.reshape(A.shape[0],-1,n).sum(1)
Out[42]: 
array([[ 2, 10,  9],
       [14,  5,  0],
       [13, 13,  1]])

Upvotes: 3

Related Questions