How to find consecutive positive, negative and zeroes in a numpy array?

Question

I use the following function in order to find the consecutive negative and positive numbers, now I also want to add a condition that gets the consecutive zeros as well.
How can I do that?

def consecutive_counts(arr):
'''
Returns number of consecutive negative and positive numbers
arr = np.array
negative = consecutive_counts()[0]
positive = consecutive_counts()[1]
'''
pos = arr > 0
# is used to Compute indices that are non-zero in the flattened version of arr
idx = np.flatnonzero(pos[1:] != pos[:-1])
count = np.concatenate(([idx[0]+1], idx[1:] - idx[:-1], [arr.size-1-idx[-1]]))
negative = count[1::2], count[::2]
positive = count[::2], count[1::2]
if arr[0] < 0:
return negative
else:
return positive

thanks :)

Answer 1

Solved it :)