Numpy mean with condition

Question

I have algorithm of calculating average speed in pure python:

speed = [...]
avg_speed = 0.0
speed_count = 0

for i in speed:
    if i > 0: # I dont need zeros
        avg_speed += i
        speed_count += 1

if speed_count == 0:
    return 0.0

return avg_speed / speed_count

Is there any way to rewrite this functions with Numpy?

Answer 1

I'm surprised no one has suggested the shortest solution:

speeds_np = np.array(speeds)

speeds_np[speeds_np>0].mean()

Explanation:

speedsNp > 0 creates a boolean array of the same size satisfying the (in)equality. If fed into speedsNp, it yields only the corresponding values of speedNp where the value of the boolean array is True. All you need to do then, is just take the mean() of the result.

Answer 2

The function numpy.average can receive a weights argument, where you can put a boolean array generated from some condition applied to the array itself - in this case, an element being greater than 0:

average_speed = numpy.average(speeds, weights=(speeds > 0))

Hope this helps