Numpy / Torch：使用维度上的索引选择元素[重复]答案

【问题标题】：Numpy/Torch : Selecting elements using indices over a dimension [duplicate]Numpy / Torch：使用维度上的索引选择元素[重复]
【发布时间】：2021-06-15 10:58:56
【问题描述】：

给定如下数组：

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 方法。

【问题讨论】：

标签： python numpy pytorch

【解决方案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))

【讨论】：

【解决方案2】：

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

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

【讨论】：

【解决方案3】：

请尝试运行以下代码。

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

【讨论】：