在 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”，或者另一个可以处理仅涉及数组的嵌套循环的库？我为一般情况寻找它，而不是针对这个特定示例的解决方法（但如果你有一个，我可以打开一个单独的问答）。

Answer 1

这基本上是Numba 背后的想法。 不如 C 快，但它可以接近...它使用 jit 编译器将 python 代码编译到机器上，并且与大多数 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

纯 Python：

%%timeit

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

每个循环 338 毫秒 ± 12.6 毫秒（平均值 ± 标准偏差，7 次运行，每个循环 1 个）

使用 Numba：

%%timeit

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

每个循环 12 毫秒 ± 534 微秒（平均值 ± 标准偏差。7 次运行，每次 100 次循环）