如何用列表理解/映射替换循环

Question

在我编写的函数中，我有：

keepList = [0]
add = 0
for j in range(1,(2*n)):
    if j%2 == 0:
        add += 1
    keepList.append((j*n+add))

其中n在函数的前面定义。当n变大时，这可能会非常慢（当我调用上述函数时会变慢）。我怎么能使用地图或列表推导来代替它，从而不会降低我的功能？

Answer 1

如前所述，不要期望使用列表推导或map会大大提高性能（如果有）。但是，使用numpy会有所帮助。确实不需要循环，因为可以计算和矢量化所有东西：

def fn(n):
    arr = np.arange(2 * n)
    arr_add = np.repeat(np.arange(n), 2)
    return arr * n + arr_add