Given indexes, get values from numpy matrix

Question

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?

user3483203 · Accepted Answer · 2021-07-23 14:16:39Z

3

This is take_along_axis.

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

answered Jul 23, 2021 at 14:16

user3483203

51.3k10 gold badges72 silver badges104 bronze badges

Sign up to request clarification or add additional context in comments.

Comments

Warren Weckesser · Accepted Answer · 2021-07-23 15:16:02Z

2

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

edited Jul 23, 2021 at 15:16

answered Jul 23, 2021 at 14:09

Warren Weckesser

116k20 gold badges207 silver badges224 bronze badges

Collectives™ on Stack Overflow

Given indexes, get values from numpy matrix

2 Answers 2

Comments

Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

2 Answers 2

Comments

Comments

Your Answer

Sign up or log in

Post as a guest

Related