Marsya Irdina
Marsya Irdina

Reputation: 3

Sum 3d numpy array [i,j] along i elements

For a numpy array of shape([N,N,3]),

[[1,2,3], [1,2,3],[1,2,3]],
[[4,5,6],[4,5,6],[4,5,6]],
[[7,8,9],[7,8,9],[7,8,9]]

How can I get a [N,3] with the j elements summed up to look like this?

[2,4,6], [8,10,12], [14,16,18]

I was thinking of reshaping it, although I don't think that is necessary. I know I want to add vertically, so thats a sum down axis=0. I am just finding trouble keeping the structure as it is.

Upvotes: 0

Views: 44

Answers (2)

Akshay Sehgal
Akshay Sehgal

Reputation: 19322

How can I get a [N,3] with the j elements summed up to look like this?

Is this what you are looking for?

j = 2

a[:,:j,:].sum(1)
array([[ 2,  4,  6],
       [ 8, 10, 12],
       [14, 16, 18]])

Upvotes: 0

Agnes Kis
Agnes Kis

Reputation: 499

You are multiplicating by two, but for sum, this is a solution:

# Convert to list
>>> lst = a.tolist()

# Sum by index
>>> sum_by_index = lambda x: [x[0][i] + x[1][i] + x[2][i] for i in range(3)]

# For multiplication
>>> mply_by_index = lambda x: [x[0][i] * 2 for i in range(3)]

# Apply
>>> list(map(sum_by_index, lst))
[[3.0, 6.0, 9.0], ...]

>>> list(map(mply_by_index, lst))
[[2.0, 4.0, 6.0], ...]

Upvotes: 1

Related Questions