python - 在 numpy 中计算累积乘积或总和时提前中断

Question

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

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

是否可以根据条件提前突破 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。