for文を使わずnumpyなどでいっぺんに処理したい

Question

下記のコードで、[np.array([np.linalg.norm(list - data)for list in List]).argmin() for data in Data]をfor文を使わずにnumpyなどでいっぺんに処理したいのですがどうすればよろしいのでしょうか?
ご回答お待ちしております。

import numpy as np

def test(List,Data):
    #ListとDataのユークリッド距離が最も小さくなるようなインデックスを返す
    return [np.array([np.linalg.norm(list - data)for list in List]).argmin() for data in Data]#ここをいっぺんに処理したい

if __name__ == '__main__': 
    LIST=np.array([[10,20,30,40],[50,60,70,80],[90,100,110,120]])
    DATA=np.array([[16,72,51,62],[1,5,100,124],[120,71,17,140]])
    A=test(LIST,DATA)
    print(A)
 

Answer 1

距離行列を作成して、argmin(axis=0) を取ればよいです。

numpy のやり方

import numpy as np
from scipy.spatial import distance_matrix

a = np.array([[10, 20, 30, 40],
              [50, 60, 70, 80],
              [90, 100, 110, 120]])
b = np.array([[16, 72, 51, 62],
              [1, 5, 100, 124],
              [120, 71, 17, 140]])

# numpy
dist_matrix = np.linalg.norm(a[:, np.newaxis] - b, axis=-1)
nearest = dist_matrix.argmin(axis=0)
print(nearest)  # [1, 1, 2]

a[:, np.newaxis]: a.shape(3, 4) -> a.shape(3, 4, 1)
a[:, np.newaxis] - b: a.shape(3, 4, 1) - b.shape(3, 4) なので、ブロードキャストで a.shape(3, 4, 4) - b.shape(3, 3, 4) となる。

scipy のやり方

scipy.spatial.distance_matrix — SciPy v1.2.1 Reference Guide でも同じことができます。

# num
# scipy
dist_matrix = distance_matrix(a, b)
nearest = dist_matrix.argmin(axis=0)
print(nearest)  # [1, 1, 2]

距離行列の可視化

import matplotlib.pyplot as plt
import seaborn as sns

# 可視化する。
plt.figure(facecolor='w', figsize=(5, 5))
ax = sns.heatmap(dist_matrix, annot=True, fmt='.1f', cmap='Reds', cbar=False)
ax.set_xlabel('B')
ax.set_ylabel('A')
plt.show()

dist_matrix [i, j] は a の i 番目の要素と b の j 番目の要素との距離を表している。
b の各要素から最も近い a の要素を調べるには、距離行列の axis=0 方向に argmin をとればよいことがわかる。

Comments

[taiki9121]

ご回答ありがとうございます!
とても分かりやすかったです^^