如何用NumPy快速定位数组累加和首次为负的索引？

当然可以用NumPy优化你的需求，不过要根据数组规模选择合适的方案：

方案1：纯NumPy向量化实现（简洁高效） #

NumPy的np.cumsum可以快速计算整个数组的前缀和，之后再找到第一个小于0的前缀和对应的索引。这种方法虽然会计算完整的前缀和，但因为np.cumsum是高度优化的向量化操作，速度远快于原生Python循环，对于大部分场景（比如你的1000元素数组）完全够用。

示例代码：

import numpy as np

a = np.random.rand(1000)
a -= .55
a[0] = 1

# 计算前缀和
prefix_sum = np.cumsum(a)
# 找到第一个前缀和小于0的索引（np.argmax会返回第一个True的位置）
first_neg_idx = np.argmax(prefix_sum < 0)

# 处理所有前缀和都非负的情况
if prefix_sum[first_neg_idx] >= 0:
    print("没有找到满足条件的索引")
else:
    print(first_neg_idx)

方案2：结合Numba实现提前终止（适合超大数组） #

如果你的数组非常大，且第一个负前缀和出现的位置很早，计算完整前缀和会有点浪费。这时候可以用Numba的JIT编译把原循环转换成机器码，既保留提前终止的逻辑，又获得接近C语言的速度。

示例代码：

import numpy as np
from numba import jit

@jit(nopython=True)
def find_first_neg_index(arr):
    current_sum = 0.0
    for idx in range(len(arr)):
        current_sum += arr[idx]
        if current_sum < 0:
            return idx
    return -1  # 无满足条件的元素时返回-1

a = np.random.rand(1000)
a -= .55
a[0] = 1

result_idx = find_first_neg_index(a)
if result_idx != -1:
    print(result_idx)
else:
    print("没有找到满足条件的索引")

总结 #

• 数组规模不大时，优先选方案1，代码简洁且足够高效；
• 数组极大且终止位置靠前时，方案2能避免不必要的计算，效率更高。

内容的提问来源于stack exchange，提问作者Andreas Schuldei