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

import numpy as np

def avg_positive_speed(speed):
    s = np.array(speed)
    positives = s > 0
    if positives.any():
        return s[positives].mean()
    else:
        return 0.


speed = [1., 2., 0., 3.]
print avg_positive_speed(speed)
# prints 2.0

print avg_positive_speed([0., 0.])
# prints 0.0