問題. 1辺が100の正方形の中の格子点にある10個の点の中から最も近い2点を見つけよ。

このような問題を解くのにScipyのKDTreeを使うと簡単にかつ高速に求めることが出来たので紹介します。

まず10個の点を乱数で発生させリストにします。

from scipy.spatial import KDTree
from random import randint

N, XYrange = 10, 10**2
points = [(randint(0,XYrange),randint(0,XYrange)) for i in range(N)] 
print(points)
# [(19, 27), (86, 35), (35, 10), (42, 40), (16, 52), (87, 3), (4, 37), (88, 34), (32, 51), (39, 59)]

次にこれをKDTreeに入力します。

tree = KDTree(points)

このKDTreeクラスに対して、点 (19,27)から近い順に2点を探すqueryを行います。

d, idx = tree.query([19,27], k=2)
print(d, idx, sep='\n')
[ 0.         18.02775638]
[0 6]

帰ってくるdは1番近い点と2番目に近い点への距離。idxはその点のindexを表します。1番近い点は当然ながら自分自身になるので、2番目に近い点が知りたい点と言うことになります。

	index	距離	座標
1番近い点	0	0	(19,27)
2番目に近い点	6	18.03	(4,37)

したがって元の問題を求めるにはすべての点に対して自分以外の一番近い点までの距離を求めてその最小値が答えとなります。

幸いqueryメソッドは点をリストで与えると答えもリストで帰ってくるので。その答えから2番目に近い点の距離とindexを抜き出します。

# ---- Query the nearest point for all points
d, idx = list(map(lambda x: x[:,1], tree.query(points, k=2))) # extract column 1 (2nd closest)
print(d, idx, sep='\n')
# ---- Find the shortest distance pair from the list
midx = d.argmin()  # p1 index for the nearest pair
p1, p2 = points[midx], points[idx[midx]]
print(f"The nearest pair: p1:{p1}-p2:{p2}, distance: {(d[midx]):.03f}")
#[18.02775638  2.23606798 23.34523506 14.86606875 16.03121954 31.01612484 18.02775638  2.23606798 10.63014581 10.63014581]
#[6 7 0 8 8 7 0 1 9 8]
#The nearest pair: p1:(86, 35)-p2:(88, 34), distance: 2.236

無事一番近い2点は (85,35)と(86,34) で距離は 2.236であることが求まりました。

（開発環境：Google Colab）

この考え方はProject Euler Problem 816: Shortest distance among pointsを解くのに役に立ちます。