1

I'm trying to write a code without loops (and preferably without lambda function) for a numpy array that calculates the sum of elements up to current element of the array, and replaces element with zero if the sum is greater than a certain value, and starts counting the sum again from the next element.

example:

np array from 1 to 10 and sum limited to 10.

I can do something like this through a loop, but how without it?

input is:

[1, 2, 3, 4, 5, 6, 7, 8, 9]

array = np.arange(1,10)
sum_arr = 0

for i in range(len(array)):
    sum_arr += array[i]
    if sum_arr > 10:
        sum_arr = 0
        array[i] = 0

output is:

[1, 2, 3, 4, 0, 6, 0, 8, 0]

Andrey Urusov
  • 19
  • 3
  • 2
    You can't avoid loops; the best you can do in `numpy` is to move the loops into compiled code. And as the answers show, the closest operation is `cumsum`, or some other form of `ufunc.accumulate`. But "resetting" the sum periodically requires some creative thinking. – hpaulj Apr 08 '23 at 19:58

3 Answers3

2

Building on top of @BENY's answer to use numpy's boolean array indexing :

array = np.arange(1, 10)

L = 10 # <- change the limit here
​
cs_func = np.frompyfunc(lambda a,b: a+b if a < L+1 else b, 2, 1)
    
array[cs_func.accumulate(array) > L] = 0

Output :

>>> print(array)
​
# array([1, 2, 3, 4, 0, 6, 0, 8, 0])

Intermediates :

>>> print(cs_func.accumulate(array))

#array([1, 3, 6, 10, 15, 6, 13, 8, 17], dtype=object)
Timeless
  • 22,580
  • 4
  • 12
  • 30
  • 2
    This use of `frompyfunc` probably isn't a lot faster than the equivalent loop, but it is a nice use of its `accumulate`, implementing a complex version of `cumsum`. – hpaulj Apr 08 '23 at 19:39
  • `%%timeit` shows that the OP's loop is about *~x4* faster than the func routine approach but the latter is a lot more fancier, I agree ;). But I still don't know why the OP are trying to avoid loops and the funny thing is they updated their question (*to exclude lambdas too*) right after my answer.. – Timeless Apr 08 '23 at 19:45
1

The numpy approach I show later in my answer is actually slower than just iterating over the array. You will get the best performance by simply using numba to pre-compile your loopy function:

import numba

@numba.njit
def andrey_numba(array, limit):
    sum_arr = 0

    for i in range(len(array)):
        sum_arr += array[i]
        if sum_arr > limit:
            sum_arr = 0
            array[i] = 0
    return array

Timeless's creative application of np.ufunc.accumulate is about as fast as your loopy approach.

enter image description here

I can't think of a way to avoid loops completely, but you could loop until the cumsum contains any elements greater than 10, and subtract the sum of the previous elements from the remaining array:

def pranav(array, limit):
    cumsum_arr = array.cumsum()

    over_limit = cumsum_arr > limit
    iters = 0 # Just to keep track of how many iterations
    while over_limit.any():
        iters += 1
        over_limit_index = over_limit.argmax() # https://stackoverflow.com/a/48360950/843953
        array[over_limit_index] = 0
        cumsum_arr[over_limit_index:] -= cumsum_arr[over_limit_index]
        over_limit = cumsum_arr > limit

    return array

Which leaves you with the desired array, in fewer iters (3 instead of 9):

array([1, 2, 3, 4, 0, 6, 0, 8, 0])

However, it actually takes more time, since you end up doing a lot more calculations in each loop. IMO this is a good illustration that not everything is made more efficient by using numpy. There are probably ways to gain performance by tweaking my numpy code, but I decided it's not worth it since numba and native-python perform well enough.

Pranav Hosangadi
  • 23,755
  • 7
  • 44
  • 70
0

With no loops, a global and built-ins only:

sum = 0
def foo(x):
    global sum
    sum += x
    if sum > 10:
        sum = 0
        return 0
    return x
print(list(map(foo, [1, 2, 3, 4, 5, 6, 7, 8, 9])))

Output:

[1, 2, 3, 4, 0, 6, 0, 8, 0]

Performance:

enter image description here

C. Pappy
  • 739
  • 4
  • 13