python - numpy.mean on varying row size

Question

The numpy mean function works perfectly fine when the dimensions are the same.

a = np.array([[1, 2], [3, 4]])
a.mean(axis=1)
array([ 1.5,  3.5])

But if I do it with varrying row size it gives an error

a = np.array([[1, 2], [3, 4, 5]])
a.mean(axis=1)
IndexError: tuple index out of range

I cannot find anything on the documentation regarding this problem. I could calculate the mean myself but I would like to use the build in function for this, seeing that it should be possible.

Answer 1

Here's an approach -

# Store length of each subarray
lens = np.array(map(len,a))

# Generate IDs based on the lengths
IDs = np.repeat(np.arange(len(lens)),lens)

# Use IDs to do bin-based summing of a elems and divide by subarray lengths
out = np.bincount(IDs,np.concatenate(a))/lens

Sample run -

In [34]: a   # Input array
Out[34]: array([[1, 2], [3, 4, 5]], dtype=object)

In [35]: lens = np.array(map(len,a))
    ...: IDs = np.repeat(np.arange(len(lens)),lens)
    ...: out = np.bincount(IDs,np.concatenate(a))/lens
    ...: 

In [36]: out  # Average output
Out[36]: array([ 1.5,  4. ])

---

Simpler alternative way using list comprehension -

In [38]: [np.mean(i) for i in a]
Out[38]: [1.5, 4.0]