CODEWITHSUNDEEP
pythonnumpy
Riccardo Bucco
Riccardo Bucco

Reputation: 15364

Given indexes, get values from numpy matrix

Let's say I have this numpy matrix:

>>> mat = np.matrix([[3,4,5,2,1], [1,2,7,6,5], [8,9,4,5,2]])
>>> mat
matrix([[3, 4, 5, 2, 1],
        [1, 2, 7, 6, 5],
        [8, 9, 4, 5, 2]])

Now let's say I have some indexes in this form:

>>> ind = np.matrix([[0,2,3], [0,4,2], [3,1,2]])
>>> ind
matrix([[0, 2, 3],
        [0, 4, 2],
        [3, 1, 2]])

What I would like to do is to get three values from each row of the matrix, specifically values at columns 0, 2, and 3 for the first row, values at columns 0, 4 and 2 for the second row, etc. This is the expected output:

matrix([[3, 5, 2],
        [1, 5, 7],
        [5, 9, 4]])

I've tried using np.take but it doesn't seem to work. Any suggestion?

Upvotes: 1

Views: 91

Answers (2)

Warren Weckesser
Warren Weckesser

Reputation: 114791

This will do it: mat[np.arange(3).reshape(-1, 1), ind]

In [245]: mat[np.arange(3).reshape(-1, 1), ind]
Out[245]: 
matrix([[3, 5, 2],
        [1, 5, 7],
        [5, 9, 4]])

(but take_along_axis in @user3483203's answer is simpler).

Upvotes: 2

user3483203
user3483203

Reputation: 51165

This is take_along_axis.

>>> np.take_along_axis(mat, ind, axis=1)
matrix([[3, 5, 2],
        [1, 5, 7],
        [5, 9, 4]])

Upvotes: 3

Related Questions