Как можно ускорить выполнение цикла for в python?

Есть вот такой вот цикл, заполняющий мне список:

        i = 0
        j = 1
        newArray = []

        for m in range(count):  
            x = ArrayOfByte[i] + (ArrayOfByte[j] * 256)
            i += 2
            j += 2
            newArray.append((x >> swift) & 255)

Элементов в конечном списке получается много, на выполнение уходит несколько секунд, что довольно много, потому что мне не один раз приходится обращаться к методу, в котором есть этот цикл. Можно ли его как-то оптимизировать? Возможно ли преобразовать этот цикл в генератор? Где-то читал, что метод append довольно медленный, но не знаю какой-либо альтернативы ему. Может, стоит использовать массивы NumPy вместо списков?


Ответы (4 шт):

Автор решения: passant

Если заранее известно значение count (а из кода следует, что это так) то разумнее сначала создать numpy-массив этого размера

newArray=np.empty(count)

, а потом в цикле просто заполнять его элементы. Обещают ускорение на порядок.

→ Ссылка
Автор решения: Vasily

Можно попробовать запустить цикл параллельно

from joblib import Parallel, delayed
def process(i):
    #ваша функция
    
results = Parallel(n_jobs=2)(delayed(process)(i) for i in range(count))
→ Ссылка
Автор решения: CrazyElf

Подготовка искусственных данных:

import numpy as np

count = 1_000_000
swift = 4
ArrayOfByte = np.random.randint(255, size=count*2)

Проверим, сколько будет работать обычный питон и сколько он насчитает:

%%time

i = 0
j = 1
newArray = []

for m in range(count):  
    x = ArrayOfByte[i] + (ArrayOfByte[j] * 256)
    i += 2
    j += 2
    newArray.append((x >> swift) & 255)
print(sum(newArray))

Вывод:

118611957
CPU times: user 1.76 s, sys: 0 ns, total: 1.76 s
Wall time: 1.76 s

А теперь векторизованный код numpy:

%%time

newArray = ((ArrayOfByte[::2] + (ArrayOfByte[1::2] * 256)) >> swift) & 255
print(newArray.sum())

Вывод:

118611957
CPU times: user 11 ms, sys: 0 ns, total: 11 ms
Wall time: 14.6 ms

Сумма чисел получилась та же, т.е. код выполняет одинаковые вычисления. При этом векторизованный код выполняется в 100 раз быстрее.

P.S. На более "мелких" типах данных вроде np.uint8 разница между питоном и numpy ещё больше - в 1000 раз. Numpy на них работает ещё быстрее, а Python ещё медленнее. По умолчанию используется тип np.int64.

→ Ссылка
Автор решения: extrn

Если ArrayOfByte – это действительно массив uint8, тогда, чтобы представить его в виде массива uint16 (а к этому и сводится ваша задача), проще и быстрее будет получить соответствующее представление массива.

>>> ArrayOfByte = np.array([1, 0, 0, 1, 255, 255], dtype='B')
>>> ArrayOfByte
array([  1,   0,   0,   1, 255, 255], dtype=uint8)
>>> ArrayOfByte.view('<u2')
array([    1,   256, 65535], dtype=uint16)

Итого

(ArrayOfByte.view('<u2') >> swift) & 255

Если на входе не массив uint8 и/или на выходе не uint16, нужно будет произвести соответствующие преобразования

(ArrayOfByte.astype('B').view('<u2') >> swift).astype('B')

Здесь на выходе массив байтов, и необходимость наложения маски 255 отпадает.

→ Ссылка