using a list comprehension to perform cumsum on multiple arrays

user_16004199 0
Viewed 58

I am trying to write a code where it checks for the np.cumsum() of both a and b separated into positive and negative values. so for the first row in the outputs [12, 101, 111] it combines all the values that are over 0 being [12, 12+89, 12+89+10]. The outputs for the first and second row are of the cumsum values that are over 0 and the values for the third and fourth row are the ones for the cumsum values that are below 0. How would i be able to add that functionality to the list comprehension below and get the Expected output below?

a = np.array([12, -5, -55, 89, 10, -5.5])
b = np.array([-5, -4.5, 12.4, 11, 16])
Sum_long_profits = [(row > 0).cumsum() for row in [a, b]]

Expected output:

[[12, 101, 111],
[12.4, 23.4, 39.4],
[-5, -60, -65.5],
[-4, -9.5]]
multidimensional-array for-loop numpy arrays python
2 Answers

One approach:

import numpy as np

a = np.array([12, -5, -55, 89, 10, -5.5])
b = np.array([-5, -4.5, 12.4, 11, 16])
sum_long_profits = [row[row > 0].cumsum() for row in [a, b]] + [row[row < 0].cumsum() for row in [a, b]]
print(sum_long_profits)

Output

[array([ 12., 101., 111.]), array([12.4, 23.4, 39.4]), array([ -5. , -60. , -65.5]), array([-5. , -9.5])]

As an alternative in just one single list comprehension, you could do:

import numpy as np
from operator import lt, gt

a = np.array([12, -5, -55, 89, 10, -5.5])
b = np.array([-5, -4.5, 12.4, 11, 16])
sum_long_profits = [row[operation(row, 0)].cumsum() for operation in [gt, lt] for row in [a, b]]
print(sum_long_profits)

Output

[array([ 12., 101., 111.]), array([12.4, 23.4, 39.4]), array([ -5. , -60. , -65.5]), array([-5. , -9.5])]
user_40015920

Just use two list comprehensions:

[arr[arr > 0].cumsum() for arr in [a, b]] + [arr[arr <= 0].cumsum() for arr in [a, b]]

Output:

[array([ 12., 101., 111.]), array([12.4, 23.4, 39.4]), array([ -5. , -60. , -65.5]), array([-5. , -9.5])]

But I don't like to write redundant code, so it's better with functions:

def func(*arrs, func):
    return [arr[func(arr)].cumsum() for arr in arrs]

print(func(a, b, func=np.int64(0).__le__) + func(a, b, func=np.int64(0).__gt__))

Output:

[array([ 12., 101., 111.]), array([12.4, 23.4, 39.4]), array([ -5. , -60. , -65.5]), array([-5. , -9.5])]
user_87083640
Related