CODEWITHSUNDEEP
pythonnumpy
MarujoRe
MarujoRe

Reputation: 202

Python Numpy optimized calculate mean of an array on indices given by other array

How can I calculate the mean values of a numpy array y using x indices from another array?

import numpy
x = numpy.array([100, 100, 20000, 20000, 100, 13, 100, numpy.nan])
y = numpy.array([10, 20, 30, 40, numpy.nan, 50, 60, 70])

expected result:

result['13']: (50)/1

result['100']: (10+20+60)/3

result['20000']: (30+40)/2

The following code works but is not efficient due to the size of my real datasets:

result = {}
unique = numpy.unique(x[~numpy.isnan(x)]).astype(int)
for elem in unique:
    pos = numpy.where(x == elem)
    avg = numpy.nanmean(y[pos])
    result[elem]=avg
print(result)

I've read about numpy.bincount, but wasn't capable of using it.

Upvotes: 0

Views: 112

Answers (1)

Paul Panzer
Paul Panzer

Reputation: 53029

Here is how to use bincount:

>>> nn=~(np.isnan(x)|np.isnan(y))
>>> xr,yr = x[nn],y[nn]
>>> unq,idx,cnt=np.unique(xr,return_inverse=True,return_counts=True)
>>> dict(zip(unq,np.bincount(idx,yr)/cnt))
{13.0: 50.0, 100.0: 30.0, 20000.0: 35.0}

Upvotes: 1

Related Questions