python - 列表推导将数组中的前n个项目替换为0

我正在编写一个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

最佳答案

总而言之：去掉对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。
从一个0数组开始，用1填充它，效果更好：

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

在我的例子中，这进一步将执行时间减少到26.3mks，速度提高了46%。

关于python - 列表推导将数组中的前n个项目替换为0，我们在Stack Overflow上找到一个类似的问题：https://stackoverflow.com/questions/50242502/