Cupy/Numpy数组索引忽略重复项：如何让重复索引生效？

问题说明 #

当用数组作为索引执行元素累加时，CuPy（及Numpy）默认仅对唯一索引执行一次操作，重复索引的累加逻辑会被忽略。比如你的示例代码：

import cupy as cp

matrix = cp.zeros((3, 3))
xi = cp.asarray([0, 1, 1, 2])
yi = cp.asarray([0, 1, 1, 2])
matrix[xi, yi] += 1
print(matrix.get())

实际输出：

[[1. 0. 0.]
 [0. 1. 0.]
 [0. 0. 1.]]

但我们期望重复的(1,1)索引能累加两次，得到：

[[1. 0. 0.]
 [0. 2. 0.]
 [0. 0. 1.]]

解决方案 #

方法1：使用cupy.add.at()（推荐） #

add.at()是CuPy提供的原地操作函数，会遍历所有索引（包括重复项）执行累加，不会忽略重复索引。修改代码如下：

import cupy as cp

matrix = cp.zeros((3, 3))
xi = cp.asarray([0, 1, 1, 2])
yi = cp.asarray([0, 1, 1, 2])
# 用add.at替代直接的+=操作
cp.add.at(matrix, (xi, yi), 1)
print(matrix.get())

运行后直接得到期望输出。

方法2：手动计算重复次数赋值 #

如果需要对重复次数做额外处理，可以先提取唯一索引和对应出现次数，再批量赋值：

import cupy as cp

matrix = cp.zeros((3, 3))
xi = cp.asarray([0, 1, 1, 2])
yi = cp.asarray([0, 1, 1, 2])

# 将二维索引转为一维扁平索引
flat_indices = xi * matrix.shape[1] + yi
# 获取唯一索引及对应的重复次数
unique_indices, counts = cp.unique(flat_indices, return_counts=True)
# 转回二维索引
u_xi = unique_indices // matrix.shape[1]
u_yi = unique_indices % matrix.shape[1]
# 赋值累加次数
matrix[u_xi, u_yi] += counts

print(matrix.get())

此方法同样能得到正确结果，适合需要对重复次数做二次加工的场景。

原理补充 #

普通索引赋值（如matrix[xi, yi] += 1）会先把索引对应元素提取为临时数组，执行加法后再赋值回原数组——这个过程中重复索引对应的元素只会被计算一次。而add.at()会直接在原数组上逐个处理所有索引项，因此能正确实现重复索引的累加逻辑。

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