Home > other >  Sum columns in numpy 2D array
Sum columns in numpy 2D array

Time:12-31

I have a 2D NumPy array V:

import numpy as np
np.random.seed(10)

V = np.random.randint(-10, 10, size=(6,8))

This gives V as:

[[ -1  -6   5 -10   7   6   7  -2]
 [ -1 -10   0  -2  -6   9   6  -6]
 [  5   1   1  -9  -2  -6   4   7]
 [  9   3  -5   3   9   3   2  -9]
 [ -6   8   3   1   0  -1   5   8]
 [  6  -3   1   7   4  -3   1  -9]]

Now, I have 2 lists, r1 and r2, containing column indices as follows:

r1 = [1, 2, 5]
r2 = [3, 4, 7]

What I want is to add the columns of V based on the indices pair (r1, r2) and store it in column indices r1. That is, for this case, add columns (1, 3), (2, 4) and (5, 7) and store them respectively in columns 1, 2 and 5 of V.

It can be done easily like this:

V[:, 1] = V[:, [1,3]].sum(axis=1)
V[:, 2] = V[:, [2,4]].sum(axis=1)
V[:, 5] = V[:, [5,7]].sum(axis=1)

which gives V as:

[[ -1 -16  12 -10   7   4   7  -2]
 [ -1 -12  -6  -2  -6   3   6  -6]
 [  5  -8  -1  -9  -2   1   4   7]
 [  9   6   4   3   9  -6   2  -9]
 [ -6   9   3   1   0   7   5   8]
 [  6   4   5   7   4 -12   1  -9]]

My concern is that is there a way we can do it without loops? Thanks in advance :)

CodePudding user response:

Just add V[:, r2] at V[:, r2], like below:

V[:, r1]  = V[:, r2]
print(V)

Output

[[ -1 -16  12 -10   7   4   7  -2]
 [ -1 -12  -6  -2  -6   3   6  -6]
 [  5  -8  -1  -9  -2   1   4   7]
 [  9   6   4   3   9  -6   2  -9]
 [ -6   9   3   1   0   7   5   8]
 [  6   4   5   7   4 -12   1  -9]]
  • Related