矢量化不同指数的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）坐标，而是index [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的总和......