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

时间:2015-12-01 18:16:09

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

1 个答案:

答案 0 :(得分:10)

我认为你可以做点什么

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!