python - 使用列表理解对多个数组执行 cumsum

Question

我正在尝试编写一个代码来检查 np.cumsum()两者 a和b分为正值和负值。所以对于输出中的第一行 [12, 101, 111]它组合了所有大于 0 的值 [12, 12+89, 12+89+10] 。第一行和第二行的输出是超过 0 的累积和值，第三行和第四行的输出是低于 0 的累积和值。我如何才能将该功能添加到列表中理解下面的内容并得到下面的预期输出？

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]]

预期输出:

[[12, 101, 111],
[12.4, 23.4, 39.4],
[-5, -60, -65.5],
[-4, -9.5]]

Answer 1

一种方法:

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)

输出

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

作为单一列表理解的替代方案，您可以这样做:

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)

输出

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