如何加快 python 陣列上的操作？

Question

我有一個形狀為（N x 1）的一維 numpy 數組A ，其中 N=4000000。 我需要在每個元素上運行 function (fcn)。

import bumpy as np
def fcn(a, B, c):
    return a - np.sqrt(c) * B) / np.sqrt(1 - c)

由於 B 是 (M=100 by P=20) 並且 c 是標量，因此 function 的 output 是一個大小為 M x P7A94F14Z 的數組。

我嘗試了不同的方法，但都很慢，例如：

for a in A:
    result = fcn(a, B, c)
    # do something with result...

您知道如何加快速度，例如通過多處理或任何其他方式嗎？

Answer 1

如果你有足夠的 memory，最簡單的是@hpaulj 建議：

for result in (A[:,None,None] - np.sqrt(c) * B) / np.sqrt(1 - c):
    # do stuff

如果您缺少 memory，請至少預先計算您能做到的：

X = 1 / np.sqrt(1 - c)
Y = np.sqrt(c) * B / np.sqrt(1 - c)
for a in A:
    result = a * X - Y
    # do stuff