Vary the elements in a numpy array

Question

I have a numpy array:

a = [3., 0., 4., 2., 0., 0., 0.]

I would like a new array, created from this, where the non zero elements are converted to their value in zeros and zero elements are converted to a single number equal to the number of consecutive zeros i.e:

b = [0., 0., 0., 1., 0., 0., 0., 0., 0., 0., 3.]

Looking for a vectorized way to do this as the array will have > 1 million elements. Any help much appreciated.

Answer 1

This should do the trick, it roughly works by 1) finding all the consecutive zeros and counting them, 2) computing the size of the output array and initializing it with zeros, 3) placing the counts from part 1 in the correct places.

def cz(a):
    a = np.asarray(a, int)

    # Find where sequences of zeros start and end
    wz = np.zeros(len(a) + 2, dtype=bool)
    wz[1:-1] = a == 0
    change = wz[1:] != wz[:-1]
    edges = np.where(change)[0]
    # Take the difference to get the number of zeros in each sequence
    consecutive_zeros = edges[1::2] - edges[::2]

    # Figure out where to put consecutive_zeros
    idx = a.cumsum()
    n = idx[-1] if len(idx) > 0 else 0
    idx = idx[edges[::2]]
    idx += np.arange(len(idx))

    # Create output array and populate with values for consecutive_zeros
    out = np.zeros(len(consecutive_zeros) + n)
    out[idx] = consecutive_zeros
    return out