Count since last occurence in NumPy

Question

Seemingly straightforward problem: I want to create an array that gives the count since the last occurence of a given condition. In this condition, let the condition be that a > 0:

in: [0, 0, 5, 0, 0, 2, 1, 0, 0]

out: [0, 0, 0, 1, 2, 0, 0, 1, 2]

I assume step one would be something like np.cumsum(a > 0), but not sure where to go from there.

Edit: Should clarify that I want to do this without iteration.

Answer 1

Numpy one-liner:

x = numpy.array([0, 0, 5, 0, 0, 2, 1, 0, 0])
result = numpy.arange(len(x)) - numpy.maximum.accumulate(numpy.arange(len(x)) * (x > 0))

Gives

[0, 1, 0, 1, 2, 0, 0, 1, 2]

If you want to have zeros in the beginning, turn it to zero explicitly:

result[:numpy.nonzero(x)[0][0]] = 0

Answer 2

根据条件拆分数组，并使用剩余片段的长度以及数组中第一个和最后一个元素的条件状态。

Answer 3

A pure python solution:

result = []
delta = 0
for val in [0, 0, 5, 0, 0, 2, 1, 0, 0]:
    delta += 1
    if val > 0:
        delta = 0            
    result.append(delta)