Python 无迭代的累积计数（单位：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])

这是因为我们可以使用累计最大值来计算“超支”：

---

健全性测试：

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!

---

干得好！仅供参考，这个问题的公认答案解决了一个几乎相同的问题，但在某些累积值为负值时也有效。

>>> test()
hooray!