List comprehension to replace the first n items in array with 0

Question

I'm writing a jitted function to replace the first N elements in a multi-dimensional array with 0. I will be doing this many, many times so speed is important. @njit speeds it up significantly but I'm wondering if there is a way to get rid of the for loop using list comprehension. Would that help improve the efficiency of this? Any suggestions?

 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

To summarize: get rid of all calls to int and len ; get rid of multiplication by 0; generate the original array efficiently.

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

This shrinks the execution time from 49 mks to 31.7 mks.

Starting with an array of zeros and stuffing it with 1s works even better:

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

In my case, this further reduces the execution time to 26.3 mks, a 46% speedup.

Answer 2

I got the run time down by about 30% using:

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

• Your version:

  41 µs ± 734 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

• This version:

  28.2 µs ± 353 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

This may just be a quirk of what random lengths were selected, but at the very least not using 0*len(lengths[i]) and either using np.ones(...) or np.zeros(...) and not 1 + np.zeros(...) is a good start.