CODEWITHSUNDEEP
pythonarraysnumpycumsum
Łukasz Szymankiewicz
Łukasz Szymankiewicz

Reputation: 137

cumulative sum in numpy array with condition for stop

I want to optimize my numpy code, Im using large arrays so efficiency is required. I tried to omit using for-looop if possible. Let`s assume simple 2-d array

1 3 5
2 0 1
5 6 2

My task is to choose this values from columns until cumsum reaches certain value (cutting values to it if needed). Lets, name this value as clip. So after this operation I`ll have array like this:

1 3 3
2 0 0
0 0 0

I get an, rather naive idea, to calculate it with simple transformations:

array_clipped = np.clip(array, 0, clip)
array_clipped_cumsum = np.cumsum(array_clipped, axis=0)
difference = clip - cumsum
difference_trimmed = np.where(difference<0, temp, 0)
final = array_clipped + difference_trimmed
final_clean = np.where(final>=0, final, 0)

As this code works, it looks very dirty and non-numpy.

Upvotes: 2

Views: 1595

Answers (2)

Paul Panzer
Paul Panzer

Reputation: 53079

Here is another one-liner:

A = np.random.randint(0,10,(6,4))
A
# array([[0, 8, 7, 6],
#        [3, 2, 0, 4],
#        [5, 6, 6, 4],
#        [4, 5, 0, 3],
#        [7, 9, 6, 8],
#        [0, 9, 8, 3]])
cap = 15

np.diff(np.minimum(A.cumsum(0),cap),axis=0,prepend=0)
# array([[0, 8, 7, 6],
#        [3, 2, 0, 4],
#        [5, 5, 6, 4],
#        [4, 0, 0, 1],
#        [3, 0, 2, 0],
#        [0, 0, 0, 0]])

Or in two lines avoiding the slow prepend:

out = np.minimum(A.cumsum(0),cap)
out[1:] -= out[:-1]
out
# array([[0, 8, 7, 6],
#        [3, 2, 0, 4],
#        [5, 5, 6, 4],
#        [4, 0, 0, 1],
#        [3, 0, 2, 0],
#        [0, 0, 0, 0]])

Upvotes: 1

Divakar
Divakar

Reputation: 221614

A cleaner way would be -

# a is input array and clip is the clipping value
c = a.cumsum(0)
out = (a-c+c.clip(max=clip)).clip(min=0)

Upvotes: 1

Related Questions