将数组中的前n项替换为0

Question

我正在编写一个jitted函数，用0替换多维数组中的前N个元素。我会做很多很多次，所以速度很重要。@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 mks缩短到31.7 mks。你知道吗

从零数组开始，用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 mks，提高了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 µs ± 734 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

• 此版本：

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

这可能只是随机长度选择的一个怪癖，但至少不使用0*len(lengths[i])，使用np.ones(...)或np.zeros(...)而不是1 + np.zeros(...)是一个好的开始。你知道吗