没有迭代的NumPy中的累积计数

我有一个像这样的数组:

a = np.array([0.1, 0.2, 1.0, 1.0, 1.0, 0.9, 0.6, 1.0, 0.0, 1.0])

我想有一个运行计数器的1.0实例是当它遇到0.0重置,这样的结果将是:

[0, 0, 1, 2, 3, 3, 3, 4, 0, 1]

我最初的想法是使用类似b = np.cumsum(a [a == 1.0])的东西,但我不知道如何(1)修改它以重置为零或(2)如何构造它所以输出数组与输入数组的形状相同.有没有想法如何做到这一点没有迭代？

我想你可以做点什么

def rcount(a):
    without_reset = (a == 1).cumsum()
    reset_at = (a == 0)
    overcount = np.maximum.accumulate(without_reset * reset_at)
    result = without_reset - overcount
    return result

这给了我

>>> a = np.array([0.1, 0.2, 1.0, 1.0, 1.0, 0.9, 0.6, 1.0, 0.0, 1.0])
>>> rcount(a)
array([0, 0, 1, 2, 3, 3, 3, 4, 0, 1])

这是有效的,因为我们可以使用累积最大值来计算"过度计数":

>>> without_reset * reset_at
array([0, 0, 0, 0, 0, 0, 0, 0, 4, 0])
>>> np.maximum.accumulate(without_reset * reset_at)
array([0, 0, 0, 0, 0, 0, 0, 0, 4, 4])

理智测试:

def manual(arr):
    out = []
    count = 0
    for x in arr:
        if x == 1:
            count += 1
        if x == 0:
            count = 0
        out.append(count)
    return out

def test():
    for w in [1, 2, 10, 10**4]:
        for trial in range(100):
            for vals in [0,1],[0,1,2]:
                b = np.random.choice(vals, size=w)
                assert (rcount(b) == manual(b)).all()
    print("hooray!")

然后

>>> test()
hooray!