накопительная сумма в массиве numpy с условием остановки

Question

Я хочу оптимизировать мой код, я использую большие массивы, поэтому требуется эффективность. Я попытался опустить использование for-looop, если это возможно. Давайте предположим, что простой двумерный массив

1 3 5
2 0 1
5 6 2

Моя задача состоит в том, чтобы выбирать эти значения из столбцов до тех пор, пока cumsum не достигнет определенного значения (при необходимости обрезая его до значений). Давайте назовем это значение как клип. Поэтому после этой операции у меня будет такой массив:

1 3 3
2 0 0
0 0 0 

У меня довольно наивная идея вычислить его с помощью простых преобразований:

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)

Поскольку этот код работает,это выглядит очень грязно и не numpy.

Answer 1

Вот еще одна строка:

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

Или в две строки, избегая медленных 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]])

Answer 2

Более чистый путь будет -

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