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

Question

我有一个这样的数组：

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

我想要一个instances of 1.0的运行计数器，遇到 0.0 时会重置，所以结果是：

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

我最初的想法是使用 b = np.cumsum(a[a==1.0]) 之类的东西，但我不知道如何 (1) 将其修改为重置为零或 (2) 完全如何构造它，使输出数组与输入数组的形状相同。任何想法如何在不迭代的情况下做到这一点？

Answer 1

我认为你可以做类似的事情

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!