Понимание списка для замены первых n элементов в массиве на 0

Question

Я пишу объединенную функцию, чтобы заменить первые N элементов в многомерном массиве на 0. Я буду делать это много, много раз, поэтому скорость важна. @njit значительно ускоряет его, но мне интересно, есть ли способ избавиться от цикла for, используя понимание списка. Поможет ли это повысить эффективность этого? Есть предложения?

 import numpy as np
 from numba import njit

 lengths=np.random.randint(0,365, size=20)

 @njit
 def availarray(lengths):
     out=1+np.zeros((365, len(lengths)))
     for i in range(int(len(lengths))):
         out[:int(lengths[i]), i]=0*int(lengths[i])
     return out

Answer 1

Подводя итог: избавьтесь от всех звонков на int и len; избавиться от умножения на 0; эффективно создать исходный массив.

def availarray(lengths):
    out = np.ones((365, lengths.size))
    for i in range(lengths.size):
        out[:lengths[i], i] = 0
    return out

Это сокращает время выполнения с 49 до 31,7 мкс.

Начиная с массива нулей и заполняя его 1с, работает еще лучше:

def availarray(lengths):
    out = np.zeros((365, lengths.size))
    for i in range(lengths.size):
        out[lengths[i]:, i] = 1
    return out

В моем случае это еще больше сокращает время выполнения до 26,3 мкс, что составляет 46% ускорения.

Answer 2

Время выполнения уменьшилось примерно на 30%, используя:

def avail_array(lengths):
    out = np.zeros((365, len(lengths)))
    for i in range(int(len(lengths))):
        out[int(lengths[i]):, i] = 1
    return out

• Ваша версия:

  41 мкс ± 734 нс на цикл (среднее ± стандартное отклонение из 7 циклов, 10000 циклов каждый)

• Эта версия:

  28,2 мкс ± 353 нс на цикл (среднее ± стандартное отклонение из 7 циклов, 10000 циклов каждый)

Это может быть просто причудой того, какие случайные длины были выбраны, но, по крайней мере, не использование 0*len(lengths[i]) и использование либо np.ones(...) или np.zeros(...), а не 1 + np.zeros(...) - хорошее начало.