How to count distance to the previous zero in pandas series?

Question

I have the following pandas series (represented as a list):

[7,2,0,3,4,2,5,0,3,4]

I would like to define a new series that returns distance to the last zero. It means that I would like to have the following output:

[1,2,0,1,2,3,4,0,1,2]

How to do it in pandas in the most efficient way?

Answer 1

The complexity is O(n). What will slow it down is doing a for loop in python. If there are k zeros in the series, and log k is negligibile comparing to the length of series, an O(n log k) solution would be:

>>> izero = np.r_[-1, (ts == 0).nonzero()[0]]  # indices of zeros
>>> idx = np.arange(len(ts))
>>> idx - izero[np.searchsorted(izero - 1, idx) - 1]
array([1, 2, 0, 1, 2, 3, 4, 0, 1, 2])