python - 向量化不同索引的 numpy 数组调用

Question

我有一个 2D numpy 数组和一个索引列表列表，我希望从 numpy 数组中计算相应的 1D 向量的总和。这可以通过 for 循环或列表理解轻松完成，但我想知道是否可以对其进行矢量化。使用类似的代码，我从矢量化中获得了大约 40 倍的加速。

这是示例代码：

import numpy as np
indices = [[1,2],[1,3],[2,0,3],[1]]
array_2d = np.array([[0.5, 1.5],[1.5,2.5],[2.5,3.5],[3.5,4.5]])
soln = [np.sum(array_2d[x], axis=-1) for x in indices]

（编辑）：请注意，对于 array_2d，索引不是 (x,y) 坐标，而是 indices[0] = [1,2] 表示 array_2d 中的第一个和第二个向量（行）。索引中每个列表的元素数量是可变的。

这是我希望能够做到的：

vectorized_soln = np.sum(array_2d[indices[:]], axis=-1)

有谁知道是否有任何方法可以实现这一目标？

Answer 1

首先，我认为您在索引的第三个元素中有错字...

最简单的方法是构建一个包含两个索引数组的 sub_array：

i = np.array([1,1,2])
j = np.array([2,3,?])
sub_arr2d = array_2d[i,j]

最后，您可以取 sub_arr2d 的总和...