Numpy.argmax() on certain axis in 2D matrix

Question

I'm currently working on creating a ranking algorithm to rank relationships between students. In my NxN matrix F: F[i, j] refers to the relationship between student i and student j. The higher the value, the stronger the relationship.

My problem is such. For creating a cluster of k students, I first choose the argmax of F. If the argmax returns index (i, j), I add students i and j to my cluster. I then want to find the argmax of F along axis i and j, and take the larger of the two as the next student in my cluster. I then repeat this process along the axis of each student in the cluster until I have k students in my cluster.

Where I am confused: numpy.argmax() takes a flattened axis as the argument for which axis to search through. How do I flatten axis i and j so that I can do numpy.argmax(F, flattenedAxis)?

I'm a relative beginner to Python, so this probably has an easy answer whose existence I was unaware of. Thanks in advance.

Answer 1

This will give you the index of the largest element in matrix F.

ind = np.unravel_index(np.argmax(F),F.shape)

This will give you the row number of the largest element in column j.

max_row_in_col_j = np.argmax(F[:,j])

This will give you the column number of the largest element in row i.

max_col_in_row_i = np.argmax(F[i,:])