python - 使用 numpy 加速 for 循环

Question

下一个 for 循环如何通过 numpy 获得加速？我想这里可以使用一些奇特的索引技巧，但我不知道是哪一个(这里可以使用 einsum 吗？)。

a=0
for i in range(len(b)):
    a+=numpy.mean(C[d,e,f+b[i]])*g[i]

编辑:C 是一个形状类似于 (20, 1600, 500) 的 numpy 3D 数组。d,e,f 是“有趣”点的索引(d,e,f 的长度相同，都在 900 左右)b 和 g 的长度相同(大约 50)。对 C 中索引为 d,e,f+b[i]

的所有点取均值

Answer 1

您可以使用以下技巧:

C[d, e][:, np.add.outer(f, b)].dot(g).diagonal().mean()

通过过早地采用将形成对角线的元素来进一步改进:

C[d, e][np.arange(len(d))[:, None], np.add.outer(f, b)].dot(g).mean()