Numpy/Torch：使用维度上的索引选择元素

Question

给定如下所示的数组：

np.arange(12).reshape(4,3)
Out[119]: 
array([[ 0,  1,  2],
       [ 3,  4,  5],
       [ 6,  7,  8],
       [ 9, 10, 11]])

我想使用索引列表[0, 2, 1, 2]从每一行中选择一个元素，以创建[0, 5, 7, 11]的4x1数组

有没有简单的方法来建立索引。我能看到的最接近的方法是pytorch中的gather方法

Answer 1

>>> import torch
>>> import numpy as np
>>> s = np.arange(12).reshape(4,3)
>>> s = torch.tensor(s)
>>> s
tensor([[ 0,  1,  2],
        [ 3,  4,  5],
        [ 6,  7,  8],
        [ 9, 10, 11]])
>>> idx = torch.tensor([0, 2, 1, 2])
>>> torch.gather(s,-1 ,idx.unsqueeze(-1))
tensor([[ 0],
        [ 5],
        [ 7],
        [11]])

torch.gather(s,-1 ,idx.unsqueeze(-1))

Answer 2

请尝试运行以下代码

import numpy as np
x = [[ 0,  1,  2],
       [ 3,  4,  5],
       [ 6,  7,  8],
       [ 9, 10, 11]]
index_array = [0, 2, 1, 2]       
index = 0
result = []
for item in x:
    result.append(item[index_array[index]])
    index += 1
print (result)

结果如下

[0, 5, 7, 11]
> 

Answer 3

arr[[0,1,2,3], [0,2,1,2]]

或者，如果您更喜欢np.arange(4)作为第一个索引数组