Numpy/Torch : Selecting elements using indices over a dimension

Question

Given an array like below:

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

I want to select a single element from each of the rows using a list of indices [0, 2, 1, 2] to create a 4x1 array of [0, 5, 7, 11].

Is there any easy way to do this indexing. The closest I could see was the gather method in pytorch.

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

Please try to run the following code.

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)

Here is the result.

[0, 5, 7, 11]
> 

Answer 3

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

or if you prefer np.arange(4) for the 1st indexing array.