2つのベクトル集合の中で距離が小さい組み合わせを求める

毎回忘れるのでメモ。

1つのベクトル集合$X$内のベクトル間距離の場合、距離行列が対称行列なので、triu_indicesで狭義上三角行列のインデックスを取得して、その部分だけargsortする。

import numpy as np
from scipy.spatial import distance_matrix

sample_num = 100
dim = 8
X = np.random.randn(sample_num, dim)
D = distance_matrix(X, X)
ti = np.triu_indices(sample_num, 1)
sorted_indices = np.argsort(D[ti])
ix, iy = ti[0][sorted_indices], ti[1][sorted_indices]
print(D[ix, iy])

2つのベクトル集合$X$, $Y$内のベクトル間距離の場合、ravelでflattenして、argsortして、unravel_indexで戻す。

import numpy as np
from scipy.spatial import distance_matrix

sample_num_x = 100
sample_num_y = 200
dim = 8
X = np.random.randn(sample_num_x, dim)
Y = np.random.randn(sample_num_y, dim)
D = distance_matrix(X, Y)
ix, iy = np.unravel_index(np.argsort(np.ravel(D)), (sample_num_x, sample_num_y))
print(D[ix, iy])

類似度の場合は$D$をマイナスにする。

Top-kだけが欲しいケースがほとんどだと思うが、その場合argpartitionしてからソートしても良い。