Question

假设我有一个范围r=numpy.array(range(1, 6))，我正在使用numpy.cumsum(r)计算累计金额。但是，由于累积结果必须小于10的条件，我希望返回[1, 3, 6, 10, 15]而不是返回[1, 3, 6]。

如果数组非常大，我希望在开始计算多余的值之前将累积和打破，稍后将丢弃。当然，为了这个问题，我在这里琐碎一切。

是否可以根据条件提前退出cumsum或cumprod？

Answer 1

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

def limited_cumsum(x, limit):
    y = []
    sm = 0
    for item in x:
        sm += item
        if sm > limit:
            return y
        y.append(sm)
    return y

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

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

Answer 2

根据您计算累积总和的数量以及预期达到目标值的速度，可能会更快地逐步计算累积总和。

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

在计算累积产品或总和的时候早点打破

2 个答案: