如何基于变长索引对Numpy数组按行求和指定元素？

问题说明 #

给定数组和对应每行的变长索引列表，需要仅对每行中指定索引位置的元素求和，输出形状与array.sum(axis=1)一致。

输入数据 #

数组：

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]
]

预期输出 #

array([19, 11, 23])

解决方案 #

1. 纯Python + NumPy快速实现 #

直接遍历每行和对应的索引，求和后转成NumPy数组，代码简洁易懂：

import numpy as np

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]
]

result = np.array([sum(row[i] for i in idx) for row, idx in zip(array, indices)])
print(result)

2. NumPy原生方法（更高效） #

如果已经将输入转为NumPy数组，用np.take方法提取指定索引元素再求和，性能更优：

import numpy as np

# 先转成NumPy数组
arr = np.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]
]

result = np.array([np.take(row, idx).sum() for row, idx in zip(arr, indices)])
print(result)

3. 掩码实现方案（满足你提到的掩码思路） #

为每行生成布尔掩码，标记需要求和的索引位置，再计算掩码后的元素和：

import numpy as np

arr = np.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]
]

result = []
for row, idx in zip(arr, indices):
    # 生成全False的掩码
    mask = np.zeros(len(row), dtype=bool)
    # 将指定索引位置设为True
    mask[idx] = True
    # 求和并加入结果列表
    result.append(row[mask].sum())

# 转成NumPy数组输出
result = np.array(result)
print(result)

内容的提问来源于stack exchange，提问作者theodosis