Question

我有一个像这样的数组：

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）相当如何结构它使输出数组与输入数组的形状相同。任何想法如何在没有迭代的情况下做到这一点？

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!

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

1 个答案: