Take values from row and column of numpy array based on another array

Question

I have a np array of weights

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

I have another numpy array containing the row and index to extract from the weight matrix.

row_col = np.array([
    [1, 1], # row 1, col 1
    [2, 2], # row 2, col 2
    [0, 2], # row 0, col 2
    [1, 0]  # row 1, col 0
])

How do I get the output:

[5, 9, 3, 4]

Mahdi F. · Accepted Answer · 2022-09-06 06:55:37Z

3

Try this:

mat[row_col[:, 0], row_col[:, 1]]

Output: array([5, 9, 3, 4])

answered Sep 6, 2022 at 6:55

Mahdi F.

24.1k5 gold badges25 silver badges32 bronze badges

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

1 Answer 1