[Python] Fastest way to find indices of highest value in a matrix iteratively and...

I'm attempting to find the "best hits" in a similarity matrix (i.e. an mxn matrix where index along each axis corresponds to the ith position in vector m and the jth position in vector n). The simplest way to explain this is finding the indices of the highest values in a matrix iteratively, excluding previously selected rows and columns. This results in min(m,n) indices chosen.

Here's my minimum reproducible example of my current implementation, using pandas:

import numpy as np
import pandas as pd

def pairwise_best_hit(sim):
xdim,ydim = np.meshgrid(np.arange(sim.shape[1]),np.arange(sim.shape[0]))
table = np.vstack((sim.ravel(),xdim.ravel(),ydim.ravel())).T
df = pd.DataFrame(table).rename(columns={0:'sim',1:'index2',2:'index1'}).sort_values('sim',ascending=False)
seq1_hits = []
seq2_hits = []
while len(df):
index1 = df.iloc[0]['index1']
index2 = df.iloc[0]['index2']
seq1_hits.append(index1)
seq2_hits.append(index2)
df = df[(df['index1']!=index1)&(df['index2']!=index2)]
return [seq1_hits,seq2_hits]


and for a matrix

sim = np.array([[1,5,6,2],[7,10,3,4],[1,5,3,7]])
pairwise_best_hit(sim)


returns

[[1, 2, 0], [1, 3, 2]]

Continue reading...