Keep an index for how many times a if statement had passed inside a list comprehension

I want for a vector with ones and zeros to check if the element in the vector is 1. If it is one I want to label this element 2 and the next one is equal to 3 etc.. This is easily done by a for loop.

x=np.array([1., 0., 0., 1., 0., 0., 1., 0., 1., 0.])

def label(vector):
U=np.copy(vector)
l=2
for i in range(len(U[:])):
    if occupied(U[i]):
        U[i]=l
        l+=1
return U

And then

label(x)

returns

array([2., 0., 0., 3., 0., 0., 4., 0., 5., 0.])

However since ths arrays become larger and eventually become matrices I thought I would be a good idea to use list comprehensions.

def count_clusters(vector):
U=np.copy(vector)
label=2
U[:] = [label, label+=1 if el==1. else el for el in U[:]]
return U

Which should look something like this. The problem is ofcourse to do the label+=1 at the same time in the list comprehension as to change the value of the element in the array. And then afterwards use the label+1 for the following label.

So I was wondering if this is possible or should I just stick to the for loop.