Populate numpy array according to predefined array values

Question

I have an A array where some indexes are populated with ones, the other values are nans.

A: [ nan nan nan nan nan nan nan 1. 1. nan nan nan nan nan nan nan nan nan nan nan 1. 1. nan nan nan nan nan nan 1. nan nan nan nan 1. nan nan nan nan nan nan nan nan nan nan nan nan nan nan nan nan nan nan nan nan nan nan nan nan nan nan nan nan nan nan nan nan nan nan nan nan nan nan]

type(A): class 'numpy.ndarray'

Is there a way to increment values after the ones (without using any kind of loops, only with filtering, indexing and broadcasting) and get this resulting B array:

B: [ nan nan nan nan nan nan nan 1. 1. 2. 3. 4. 5. 6. 7. 8. 9. 10. 11. 12. 1. 1. 2. 3. 4. 5. 6. 7. 1. 2. 3. 4. 5. 1. 2. 3. 4. 5. 6. 7. 8. 9. 10. 11. 12. 13. 14. 15. 16. 17. 18. 19 20. 21. 22. 23. 24. 25. 26. 27. 28. 29. 30. 31. 32. 33. 34. 35. 36. 37. 38. 39.]

Answer 1

Yep:

B = np.ones_like(A)
ons, = np.where(A==1)
if ons.size:
    B[:ons[0]] = np.nan
    if ons.size > 1:
        B[ons[1:]] -= np.diff(ons)
    B[ons[0]:] = B[ons[0]:].cumsum()
B
# array([nan, nan, nan, nan, nan, nan, nan,  1.,  1.,  2.,  3.,  4.,  5.,
#         6.,  7.,  8.,  9., 10., 11., 12.,  1.,  1.,  2.,  3.,  4.,  5.,
#         6.,  7.,  1.,  2.,  3.,  4.,  5.,  1.,  2.,  3.,  4.,  5.,  6.,
#         7.,  8.,  9., 10., 11., 12., 13., 14., 15., 16., 17., 18., 19.,
#        20., 21., 22., 23., 24., 25., 26., 27., 28., 29., 30., 31., 32.,
#        33., 34., 35., 36., 37., 38., 39.])