Numpy / PyTorch - 如何用不同维度的索引分配值？

Question

假设我有一个矩阵和一些索引

a = np.array([[1, 2, 3], [4, 5, 6]])
a_indices = np.array([[0,2], [1,2]])

有没有什么有效的方法来实现以下操作？

for i in range(2):
    a[i, a_indices[i]] = 100

# a: np.array([[100, 2, 100], [4, 100, 100]])

Answer 1

使用np.put_along_axis-

In [111]: np.put_along_axis(a,a_indices,100,axis=1)

In [112]: a
Out[112]: 
array([[100,   2, 100],
       [  4, 100, 100]])

或者，如果您想使用显式方式，即基于整数的索引 -

In [115]: a[np.arange(len(a_indices))[:,None], a_indices] = 100