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

- Yunqiu Xu

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

- Divakar

Answer 2

由于这个问题标记为PyTorch，所以为了完整起见，这里提供一个 PyTorch 的等效解决方案。

# make a copy of the inputs from numpy array
In [188]: at = torch.tensor(a)
In [189]: at_idxs = torch.tensor(a_indices)

我们将使用tensor.scatter_(...)进行就地替换。因此，让我们首先准备输入内容。 torch.scatter() API期望替换值（这里是100）为一个张量，并且与索引张量的形状相同。因此，我们需要创建一个填充有值100并且形状为(2, 2)的张量，因为索引张量at_idxs的形状是那样的。所以，

In [190]: replace_val = 100 * torch.ones(at_idxs.shape, dtype=torch.long)    
In [191]: replace_val 
Out[191]: 
tensor([[100, 100],
        [100, 100]])

现在执行就地替换：

# fill the values along dimension 1
In [198]: at.scatter_(1, at_idxs, replace_val)
Out[198]: 
tensor([[100,   2, 100],
        [  4, 100, 100]])

- kmario23