Python 计算累积乘积或总和（单位：numpy）时提前中断

假设我有一个range

r=numpy.array（range（1，6））

并且我正在使用

numpy.cumsum（r）

计算累积和。但不是返回

[1,3,6,10,15]

我希望它返回

[1,3,6]

，因为累积结果必须小于10

如果数组非常大，我希望在开始计算冗余值并在以后丢弃这些值之前，先将累积和取出来。当然，为了这个问题，我把这里的一切都琐碎化了

---

是否有可能根据条件提前中断

cumsum

或

cumprod

？

我认为numpy中的任何函数都不可能实现这一点，因为在大多数情况下，这些函数用于固定长度数组上的矢量化计算。一个显而易见的方法是打破Python中的标准for循环（我假设您知道）：

但这显然要比numpy的cumsum慢一个数量级

---

由于您可能需要一些非常专门的函数，因此在numpy中查找所需的确切函数所需的更改很少。您可能应该看看，它允许您实现与Python函数一样灵活的自定义函数（并且使用几乎是Python的语法），速度接近C。

根据计算累积和的数组大小以及预期达到目标值的速度，分步计算累积和可能会更快

import numpy as np

size = 1000000
target = size
def stepped_cumsum():
    arr = np.arange(size)
    out = np.empty(len(arr), dtype=int) 
    step = 1000
    last_value = 0
    for i in range(0, len(arr), step):
        np.cumsum(arr[i:i+step], out=out[i:i+step])
        out[i:i+step] += last_value
        last_value = out[i+step-1]
        if last_value >= target:
            break
    else:
        return out
    greater_than_target_index = i + (out[i:i+step] >= target).argmax()
    # .copy() required so rest of backing array can be freed
    return out[:greater_than_target_index].copy()

def normal_cumsum():
    arr = np.arange(size)
    out = np.cumsum(arr)
    return out

stepped_result = stepped_cumsum()
normal_result = normal_cumsum()
assert (stepped_result < target).all()
assert (stepped_result == normal_result[:len(stepped_result)]).all()

import numpy as np

size = 1000000
target = size
def stepped_cumsum():
    arr = np.arange(size)
    out = np.empty(len(arr), dtype=int) 
    step = 1000
    last_value = 0
    for i in range(0, len(arr), step):
        np.cumsum(arr[i:i+step], out=out[i:i+step])
        out[i:i+step] += last_value
        last_value = out[i+step-1]
        if last_value >= target:
            break
    else:
        return out
    greater_than_target_index = i + (out[i:i+step] >= target).argmax()
    # .copy() required so rest of backing array can be freed
    return out[:greater_than_target_index].copy()

def normal_cumsum():
    arr = np.arange(size)
    out = np.cumsum(arr)
    return out

stepped_result = stepped_cumsum()
normal_result = normal_cumsum()
assert (stepped_result < target).all()
assert (stepped_result == normal_result[:len(stepped_result)]).all()

In [60]: %timeit cumsum.stepped_cumsum()
1000 loops, best of 3: 1.22 ms per loop

In [61]: %timeit cumsum.normal_cumsum()
100 loops, best of 3: 3.69 ms per loop