I have a 2D matrix, where I wish to find the topK (K is some value) values in the matrix along the rows and coloumns. What I want to ask, has two parts:

Question 1:

I have a working code snippet, that uses np.argpartition and np.take_along_axis to find the required values, separately, for the rows and columns:

import numpy as np


matrix = np.array([
    [10.5, 1, 2.7],
    [8.9, 4, 3],
    [6.9, 0, 4.2],
    [11.1, 7, 0],
    [7, 3, 3.3]
])  # 5 rows, 3 columns

print(matrix.shape)

top2_cols = 2
top3_rows = 3

top2_cols_idx = np.argpartition(a=matrix, kth=-top2_cols, axis=1)[:, -top2_cols:]
matrix_top2_cols = np.take_along_axis(matrix, top2_cols_idx, axis=1)

print(matrix_top2_cols)


top3_rows_idx = np.argpartition(a=matrix, kth=-top3_rows, axis=0)[-top3_rows:, :]
matrix_top3_rows = np.take_along_axis(matrix, top3_rows_idx, axis=0)

print(matrix_top3_rows)

1. Is there a more efficient way to achieve this?
2. Is there an efficient way to achieve this while preserving the relative order

Question 2:

1. Is there a way to get a submatrix that contains the topK overall values of shape (top3_rows, top2_cols), directly?
2. If answer to above is yes, how to maintain the info of indices relative to original matrix?