如何在Numpy中基于复合主键实现Upsert（合并时累加重叠列）

给定两个3列的Numpy数组，前两列作为复合主键，需要合并数组并在主键重叠时将第三列的值相加，以下是纯Numpy实现方案：

问题复现 #

import numpy as np

x = np.array([[1, 2, 1.5], [3, 4, 2.5], [5, 6, 3.5]])
y = np.array([[1, 2, 4.5], [3, 4, 5.5], [7, 8, 6.5]])

# 期望输出
# array([[1. , 2. , 6. ],
#        [3. , 4. , 7. ],
#        [5. , 6. , 3.5],
#        [7. , 8. , 6.5]])

实现代码 #

# 合并两个数组的复合主键部分
combined_keys = np.vstack([x[:, :2], y[:, :2]])
# 获取唯一主键及原主键对应的唯一索引
unique_keys, idx = np.unique(combined_keys, axis=0, return_inverse=True)
# 合并需要求和的第三列值
values = np.hstack([x[:, 2], y[:, 2]])
# 按索引分组求和
sum_values = np.bincount(idx, weights=values)
# 拼接最终结果数组
z = np.hstack([unique_keys, sum_values.reshape(-1, 1)])

print(z)

代码解释 #

1. 合并主键集合：用np.vstack将两个数组的前两列（复合主键）合并，得到所有出现过的主键。
2. 提取唯一主键：np.unique(..., axis=0)找出所有不重复的复合主键，return_inverse=True返回每个原主键在唯一主键数组中的位置索引，用于后续分组。
3. 合并待求和值：用np.hstack将两个数组的第三列合并为一维数组。
4. 分组求和：np.bincount根据索引对值进行累加，自动处理重叠主键的求和逻辑。
5. 拼接结果：将唯一主键数组与求和后的第三列拼接成最终的3列数组。

该方案适用于任意数值类型的复合主键，且完全基于Numpy原生函数实现，无需依赖其他库。

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