Блог

У меня есть массив, и я хотел бы получить совокупную сумму, устанавливающую нижнюю границу (lb = 0), потому что в моем массиве у меня есть отрицательные элементы. Возможно ли его векторизовать? Я пробовал с циклом, numba.njit но выполнение выполняется медленнее, чем чистый Python. Ниже пример того, что я хотел бы получить.

Пример массива:

 a = numpy.array([1,1,-1,-1,-1,1,1])
  

Что я получаю с numpy.cumsum :

 [ 1  2  1  0 -1  0  1]
  

Чего я хочу:

 [ 1  2  1  0  0  1  2]
  

Функция с циклом:

 @numba.njit
def cumsum(array, lb=0):
    result = numpy.zeros(array.size)
    result[0] = array[0]
    for k in range(1, array.size):
        result[k] = max(lb, result[k-1] array[k])
    return result
  

Когда я выполняю вышеуказанное, Numba работает намного быстрее, чем чистый Python. Обратите внимание, что первый вызов cumsum() (при оформлении с @numba.njit) помощью) также выполняет компиляцию, поэтому, честно говоря, вам не следует засекать время этого первого вызова. Второй (и третий и …) вызов cumsum() будет быстрым.