numpy 如何根据可变长度索引在每行中只包含某些值,从而沿行轴对数字求和?

x33g5p2x  于 2022-11-10  发布在  其他
关注(0)|答案(1)|浏览(116)

对于以下阵列:

array = [
  [1, 5, 6, 8, 10, 3],
  [3, 2, 4, 9, 11, 7],
  [8, 0, 9, 6, 23, 4]
]

我们如何对这些指数所示的元素(每行)求和:

indices = [
  [2, 4, 5],
  [1, 3],
  [4]
]

也就是说:

  • 对于第一行,求和时只考虑索引[2, 4, 5]上的值->(6 + 10 + 3)
  • 对于第二行,汇总时只考虑索引[1, 3]上的值->(2 + 9)
  • 以此类推

产出:

array([19, 11, 23])

输出的形状与我们使用的array.sum(axis=1)相同,但并不包含每个值。取而代之的是,每行的参与者由指数数组确定。
我曾想过为此创建一个面具,但我不知道如何将索引传递给它。

ymzxtsji

ymzxtsji1#

试试这个:

arr = np.array(array)
out = np.array([arr[idx, ind].sum() for idx, ind in enumerate(indices)])
out

输出:array([19, 11, 23])

相关问题