Python에서 배열을 사용한 더 빠른 for 루프

Question

N, M = 1000, 4000000
a = np.random.uniform(0, 1, (N, M))
k = np.random.randint(0, N, (N, M))

out = np.zeros((N, M))
for i in range(N):
    for j in range(M):
        out[k[i, j], j] += a[i, j]

저는 매우 긴 for 루프로 작업합니다. %%timeit위 pass의 작업 수율 교체로

1min 19s ± 663 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

이것은 문맥 상 허용되지 않습니다 (C ++는 6.5 초 걸렸습니다). 위의 작업을 Python 객체로 수행 할 이유가 없습니다. 배열에는 잘 정의 된 유형이 있습니다. 이것을 C / C ++에서 확장으로 구현하는 것은 개발자와 사용자 모두에게 과도한 일입니다. 나는 루프에 배열을 전달하고 산술을 수행합니다.

Numpy에게 "이 논리를 C로 이동"이라고 알려주는 방법이 있습니까? 아니면 배열 만 포함하는 중첩 루프를 처리 할 수있는 다른 라이브러리가 있습니까? 이 특정 예제에 대한 해결 방법이 아닌 일반적인 경우를 찾습니다 (하지만있는 경우 별도의 Q & A를 열 수 있음).

Answer 1

이것은 기본적으로 Numba 이면의 아이디어 입니다. C만큼 빠르지는 않지만 가까이 갈 수 있습니다. jit 컴파일러를 사용하여 파이썬 코드를 기계로 컴파일하며 대부분의 Numpy 함수와 호환됩니다. (문서에서 모든 세부 정보를 찾을 수 있습니다)

import numpy as np
from numba import njit


@njit
def f(N, M):
    a = np.random.uniform(0, 1, (N, M))
    k = np.random.randint(0, N, (N, M))

    out = np.zeros((N, M))
    for i in range(N):
        for j in range(M):
            out[k[i, j], j] += a[i, j]
    return out


def f_python(N, M):
    a = np.random.uniform(0, 1, (N, M))
    k = np.random.randint(0, N, (N, M))

    out = np.zeros((N, M))
    for i in range(N):
        for j in range(M):
            out[k[i, j], j] += a[i, j]
    return out

순수 파이썬 :

%%timeit

N, M = 100, 4000
f_python(M, N)

루프 당 338ms ± 12.6ms (7 회 실행의 평균 ± 표준 편차, 각 1 회 루프)

Numba 사용 :

%%timeit

N, M = 100, 4000
f(M, N)

루프 당 12ms ± 534µs (7 회 실행의 평균 ± 표준 편차, 각 100 개 루프)