Python中按索引的j值对对应numpy数组元素求和

问题描述 #

我有两个列表Ii01和Iv01，其中Ii01包含一个存储索引的numpy数组，Iv01包含与这些索引对应的值（比如Iv01中的[2.1]对应Ii01中的[3,1]）。需要对具有相同j值（即索引数组中每个元素的第二个数值）的索引对应的元素求和：例如j=3的索引[0,3]、[2,3]、[4,3]对应的Iv01值需求和；j=4的索引[0,4]、[2,4]对应的Iv01值也需求和。

示例代码：

import numpy as np
Ii01 = [np.array([[3, 1],
       [0, 2],
       [0, 3],
       [2, 3],
       [4, 3],
       [0, 4],
       [2, 4]])]

Iv01 = [np.array([[2.1],
       [3.4],
       [1.5],
       [9.7],
       [6.5],
       [4.2],
       [1.7]])]

预期输出：

[np.array([[2.1],
       [3.4],
       [17.7],  # 1.5+9.7+6.5
       [5.9]])]  # 4.2+1.7

解决方案 #

可以利用numpy的unique函数获取唯一j值，结合where和sum实现分组计算：

import numpy as np

Ii01 = [np.array([[3, 1],
       [0, 2],
       [0, 3],
       [2, 3],
       [4, 3],
       [0, 4],
       [2, 4]])]

Iv01 = [np.array([[2.1],
       [3.4],
       [1.5],
       [9.7],
       [6.5],
       [4.2],
       [1.7]])]

# 提取所有索引的j值（数组第二列）
j_values = Ii01[0][:, 1]
# 获取唯一j值及原索引对应关系
unique_j, indices = np.unique(j_values, return_inverse=True)

# 按j值分组求和并整理为列向量格式
summed_values = np.array([Iv01[0][indices == i].sum() for i in range(len(unique_j))]).reshape(-1, 1)

# 组装成目标列表结构
result = [summed_values]
print(result)

执行后输出结果：

[array([[ 2.1],
       [ 3.4],
       [17.7],
       [ 5.9]])]

代码说明 #

• j_values = Ii01[0][:, 1]：从索引数组中提取所有j值（每个子数组的第二个元素）。
• np.unique(j_values, return_inverse=True)：获取去重后的j值，同时得到原数组每个j值在去重列表中的位置索引，用于后续分组。
• 列表推导式遍历每个唯一j值，筛选对应分组的数值求和，最后用reshape(-1,1)保持和原数组一致的列向量结构。

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