重新启动累计并获得索引（如果累计超过值）

Question

说我有一段距离x=[1,2,1,3,3,2,1,5,1,1]。

我想从x到达总和达到10的索引，在这种情况下，idx = [4,9]。

因此，满足条件后，cumsum重新启动。

我可以使用循环来完成此操作，但是对于大型阵列而言，循环速度很慢，我想知道是否可以用某种vectorized方式来执行。

Answer 1

这是一个带有numba和数组初始化的代码-

from numba import njit

@njit
def cumsum_breach_numba2(x, target, result):
    total = 0
    iterID = 0
    for i,x_i in enumerate(x):
        total += x_i
        if total >= target:
            result[iterID] = i
            iterID += 1
            total = 0
    return iterID

def cumsum_breach_array_init(x, target):
    x = np.asarray(x)
    result = np.empty(len(x),dtype=np.uint64)
    idx = cumsum_breach_numba2(x, target, result)
    return result[:idx]

时机

包括@piRSquared's solutions并使用同一篇文章中的基准测试设置-

In [58]: np.random.seed([3, 1415])
    ...: x = np.random.randint(100, size=1000000).tolist()

# @piRSquared soln1
In [59]: %timeit list(cumsum_breach(x, 10))
10 loops, best of 3: 73.2 ms per loop

# @piRSquared soln2
In [60]: %timeit cumsum_breach_numba(np.asarray(x), 10)
10 loops, best of 3: 69.2 ms per loop

# From this post
In [61]: %timeit cumsum_breach_array_init(x, 10)
10 loops, best of 3: 39.1 ms per loop

Numba：追加与数组初始化

为了更仔细地了解数组初始化有何帮助，这似乎是两个numba实现之间的最大差异，让我们将它们放在数组数据上，因为数组数据的创建本身就很耗时，而且它们都依赖于它-

In [62]: x = np.array(x)

In [63]: %timeit cumsum_breach_numba(x, 10)# with appending
10 loops, best of 3: 31.5 ms per loop

In [64]: %timeit cumsum_breach_array_init(x, 10)
1000 loops, best of 3: 1.8 ms per loop

为了强制输出拥有自己的存储空间，我们可以制作一个副本。虽然不会大幅度改变事情-

In [65]: %timeit cumsum_breach_array_init(x, 10).copy()
100 loops, best of 3: 2.67 ms per loop