Compute the mean of a two dimensional np array row wise up to a certain different element

user_13513700 0
Viewed 30

Let's consider a two dimensional array

arr = np.array([[ 0,  1,  2,  3,  4],
                [ 5,  6,  7,  8,  9],
                [10, 11, 12, 13, 14],
                [15, 16, 17, 18, 19],
                [20, 21, 22, 23, 24]])

and let's also consider another array containing the index limits.

limits = np.array([2, 3, 2, 4, 3])

Is there a way to compute the mean (row-wise) of the previous array, but considering only a bunch (given by the second array) of elements for each row? In particular, in my previous example I want the computation of the mean to take in consideration just the first 2 elements for the first row, the first 3 elements for the second row, the first 2 elements for the third row and so on so forth.

I've come up with a very basic solution, using a for loop:

arr_means = np.zeros(5, dtype=np.float32)
for i, lim in zip(range(5), limits):
    arr_means[i] = arr[i, :lim].mean()

which returns the following:

np.array([ 0.5,  6. , 10.5, 16.5, 21. ], dtype=float32)

Is it possible to vectorize this piece of code?

numpy arrays python
1 Answers

You can create a mask for the position to be averaged:

mask = np.arange(arr.shape[1])<limits[...,None]

out = (arr*mask).sum(1)/mask.sum(1)

Output:

array([ 0.5,  6. , 10.5, 16.5, 21. ])
user_42384080
Related