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] 

CodePudding user response：

Try this:

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

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