python - 如何加速python中多个数组的计算

Question

假设我们有a, b, c三个相同维度的数组，我们需要根据每一列进行计算a、b、c。

一个例子是这样的:

import numpy as np 

col = 10
row = 1000000
a = np.random.normal(size=(row, col))
b = np.random.normal(size=(row, col))
c = np.random.normal(size=(row, col))

def my_func(a, b, c):
    if a[0] + b[0] + b[-1] > c[0]:
        return a * b * c
    else:
        return a * (b[1] + b[-1]) + c[-1]

for i in range(num):
    my_func(a[:, i], b[:, i], c[:, i])

第一个想法是使用 numpy，但是由于计算不是同质的并且是在多个数组上进行的，numpy.apply_along_axes 不起作用。
在这种情况下，是否有任何建议可以在 python 中加速上述计算而不是使用 c？

Answer 1

使用 numpy 时，您需要考虑对 my_func 中的计算进行矢量化，以加快代码速度。在这种情况下，您可以尝试这样的操作:

cond = np.broadcast_to(a[0] + b[0] + b[-1] > c[0], a.shape)
result = np.where(cond, a * b * c, a * (b[1] + b[-1]) + c[-1])

在这里，像 a[0] + b[0] 这样的计算是在两个数组上执行的，而不是逐个元素执行的。