减少 numpy 数组的第三维并对值求和

Question

I think this is straightforward but I can't quite get it. I have a large 3d array and I want to reduce the 3rd dim by some factor and then sum the values to get to that reduced size. An example that works to get what I want is:

import numpy as np

arr=np.ones((10,10,16))
processed_data=np.zeros((arr.shape[0], arr.shape[1]), dtype='object')
factor=2

for i in range(arr.shape[0]):
    for j in range(arr.shape[1]):
        processed_data[i][j]=arr[i][j].reshape(int(arr.shape[2]/factor),-1).sum(axis=1)

So we take the last dimension, reshape it to an extra dimension and then sum along that dimension. In the example above the data is a 10x10x16 array of all 1s so with a factor=2 we get a 10x10x8 array out with the data all being 2s. I hope this illustrates what I am trying to achieve. If the factor would change to 4 we would get a 10x10x4 array out.

This method is not ideal as it involves creating a separate processed_data 'object' array where I would rather leave it as a 3D array, just with a reduced third dimension. It also involves iterating over every element in the 2D array which I don't think is neccessary. And it's really slow.

Any help appreciated - I suspect it is a combination of reshaping and transposing but cannot get my head around it.

Thanks.

Answer 1

I think you can reshape on the whole data and sum:

arr.reshape(*arr.shape[:2], -1, 2).sum(axis=-1)