CODEWITHSUNDEEP
pythonnumpy
Andreas Schuldei
Andreas Schuldei

Reputation: 405

Can numpy help me quickly find the index of an array, at which its sum is negativ for the first time?

I need to do something like this

import numpy as np

a = np.random.rand(1000)
a -= .55
a[0] = 1
b = 0
for i in range(len(a)):
    b += a[i]
    if b < 0:
        print(i)
        break

a lot, and preferably it should be swift. Can NumPy help me with that? I understand that NumPy is built for vector calculation and not for this. Still, it can calculate the sum of an array blazingly fast - can I give it this specific condition (the sum is negative for the first time) to stop and tell me the index number?

Upvotes: 0

Views: 68

Answers (1)

I&#39;mahdi
I&#39;mahdi

Reputation: 24059

You can use numpy.cumsum() and numpy.argmax(). First, compute the cumulative sum of the elements. Then return True/False that elements < 0 and return the first index that is True for element < 0 with argmax().

>>> (a.cumsum() < 0).argmax()

check both codes:

import numpy as np

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

def check_1(a):
    b = 0
    for i in range(len(a)):
        b += a[i]
        if b < 0:
            print(i)
            break
            
def check_2(a):
    return (a.cumsum() < 0).argmax()

Output: (Generate base random input)

>>> check_1(a)
6

>>> check_2(a)
6

Upvotes: 1

Related Questions