Fill up contiguous sequences of zeros with incrementing sequences

Is it possible to fill up zero values with arange without iterations:

Initial array:

array([128,   0, 128,   0,   0,   0, 128,   0,   0,   0,   0,   0,   0,
         0,   0, 128, 128,   0,   0, 128, 128, 128,   0,   0,   0,   0,
         0,   0,   0, 128,   0,   0,   0,   0,   0, 128,   0,   0, 128,
       128,   0,   0, 128,   0,   0,   0, 128,   0,   0, 128, 128,   0,
         0,   0, 128, 128], dtype=uint8)

new array:

array([128,   0, 128,   0,   1,   2, 128,   0,   1,   2,   3,   4,   5,
         6,   7, 128, 128,   0,   1, 128, 128, 128,   0,   1,   2,   3,
         4,   5,   6, 128,   0,   1,   2,   3,   4, 128,   0,   1, 128,
       128,   0,   1, 128,   0,   1,   2, 128,   0,   1, 128, 128,   0,
         1,   2, 128, 128], dtype=uint8)

>Solution :

This works:

zeros = a == 0
idxs = np.arange(len(a))
zero_block_starts = np.diff(zeros, prepend = 0) == 1
ranges = idxs - np.maximum.accumulate(zero_block_starts * idxs)
result = a.copy()
result[zeros] = ranges[zeros]

Result:

>>> a
array([128,   0, 128,   0,   0,   0, 128,   0,   0,   0,   0,   0,   0,
         0,   0, 128, 128,   0,   0, 128, 128, 128,   0,   0,   0,   0,
         0,   0,   0, 128,   0,   0,   0,   0,   0, 128,   0,   0, 128,
       128,   0,   0, 128,   0,   0,   0, 128,   0,   0, 128, 128,   0,
         0,   0, 128, 128], dtype=uint8)

>>> result
array([128,   0, 128,   0,   1,   2, 128,   0,   1,   2,   3,   4,   5,
         6,   7, 128, 128,   0,   1, 128, 128, 128,   0,   1,   2,   3,
         4,   5,   6, 128,   0,   1,   2,   3,   4, 128,   0,   1, 128,
       128,   0,   1, 128,   0,   1,   2, 128,   0,   1, 128, 128,   0,
         1,   2, 128, 128], dtype=uint8)

I’d be wary to use this one a uint8 dtype array though, as gaps larger than 256 elements will overflow.