[python] 二次元分布データからクエリに近しいデータを検索する方法

Question

二次元分布データからクエリに近しいデータを検索する方法を探しています。
pythonで実装したいです。

具体的には
x=[10,-1,9,12,5,6,....]
y=[0,3,-6,10,22,-4,....]
という二次元分布データにたいして、[10,8]というようなクエリとなるベクトルを与えます。
二次元分布データからクエリに近しい要素(上位５つ)を出力したいです。

forでx,yの全要素とクエリのユークリッド距離をとる方法も考えられましたが、もっとスマートな関数や方法があるのではないかと思い質問しました。

よろしくお願いします。

Answer 1

scipy.spatial.cKDTreeが便利です。
scipy.spatial.cKDTree — SciPy Reference Guide

import numpy as np
from scipy.spatial.ckdtree import cKDTree

x = [10, -1, 9, 12, 5, 6]
y = [0, 3, -6, 10, 22, -4]

xy_array = np.c_[x, y]
test_array = [[10, 8], [0, 0]]  # 近傍点を求めたい点(x, y)の配列
n = 5  # 近傍点の数
distance, index = cKDTree(xy_array).query(test_array, n)
xy_array[index]
# array([[[12, 10],    # (10, 8)に近い点5つ
#         [10,  0],
#         [-1,  3],
#         [ 6, -4],
#         [ 9, -6]],
# 
#        [[-1,  3],    # (0, 0)に近い点5つ
#         [ 6, -4],
#         [10,  0],
#         [ 9, -6],
#         [12, 10]]])

Answer 2

あんまりスマートではないですが・・・。numpyやpandas使ってやろうかと思いましたが大差無かったので、リストのままで作ってみました。

x=[10,-1,9,12,5,6]
y=[0,3,-6,10,22,-4]

def norm2(v, w):
    return (v[0]-w[0])**2+(v[1]-w[1])**2

vectors = [[s, t] for s, t in zip(x, y)]
p = [10, 8]
result = sorted(vectors, key=lambda v: norm2(v, p))[:5]
print(result)
#[[12, 10], [10, 0], [-1, 3], [6, -4], [9, -6]]