partial cumulative sum in python-CodePudding

Suppose I have a numpy array (or pandas Series if it makes it any easier), which looks like this:

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

I want to transform into an array

bar = np.array([0, 1, 2, 3, 4,0, 1, 2, 0, 1, 2, 3])

where the entry is how many steps you need to walk to the left to find a 1 in foo.

Now, obviously one can write a loop to compute bar from foo, but this will be bog slow. Is there anything more clever one can do?

CodePudding user response：

You could do cumcount with pandas

s = pd.Series(foo)
bar = s.groupby(s.cumsum()).cumcount().to_numpy()
Out[13]: array([0, 1, 2, 3, 4, 0, 1, 2, 0, 1, 2, 3], dtype=int64)

CodePudding user response：

One option, specifically for the shared example, with numpy:

# get positions where value is 1
pos = foo.nonzero()[0]
# need this when computing the cumsum
values = np.diff(pos) - 1
arr = np.ones(foo.size, dtype=int)
arr[0] = 0
arr[pos[1:]] = -values
arr.cumsum()

array([0, 1, 2, 3, 4, 0, 1, 2, 0, 1, 2, 3])