CODEWITHSUNDEEP
pythonnumpy
Keyur Potdar
Keyur Potdar

Reputation: 7238

Sum of specific group of columns for each row of a numpy array

I have the following NumPy array:

>>> a = np.arange(21).reshape(3, 7)
>>> a
array([[ 0,  1,  2,  3,  4,  5,  6],
       [ 7,  8,  9, 10, 11, 12, 13],
       [14, 15, 16, 17, 18, 19, 20]])

I want to create a new array by replacing some columns with their sums. The columns that I need to sum over will always be grouped together.

For example, keep the first 2 columns as is. Replace the column with index 2 by the sum of columns with indices 2, 3, and 4. Replace the column with index 3 by the sum of columns with indices 5 and 6.

Required output:

array([[ 0,  1,  9, 11],
       [ 7,  8, 30, 25],
       [14, 15, 51, 39]])

What I tried:

b = np.concatenate([
    a[:, :2], 
    a[:, 2:5].sum(axis=1, keepdims=True), 
    a[:, 5:].sum(axis=1, keepdims=True)
], axis=1)

This gives the required output, but I was wondering if there is a better/concise way to do this.

Upvotes: 4

Views: 715

Answers (1)

Divakar
Divakar

Reputation: 221574

Define the bins of the intervals to be summed and then use np.add.reduceat along axis=1 (for each row) -

In [37]: a
Out[37]: 
array([[ 0,  1,  2,  3,  4,  5,  6],
       [ 7,  8,  9, 10, 11, 12, 13],
       [14, 15, 16, 17, 18, 19, 20]])

In [38]: bins = [0,1,2,5]

In [39]: np.add.reduceat(a,bins,axis=1)
Out[39]: 
array([[ 0,  1,  9, 11],
       [ 7,  8, 30, 25],
       [14, 15, 51, 39]])

Upvotes: 4

Related Questions