Как заменить цикл с пониманием списка / картой

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