サブロウ丸

主にプログラミングと数学

【Python】scipy KD treeの使い方

kd treeは k dimensional treeで, k次元領域の点探索などに用いられるデータ構造です。 kd treeを取り扱うモジュールがscipyにあります。

import scipy.spatial as ss                                      
from random import random
# データ数
N = 10000
# (x座標, y座標)のデータリスト
data = [(random()*100, random()*100) for _ in range(N)]
# kd tree の作成 (leafsizeは展開をしない節内点数上限)
tree = ss.KDTree(data, leafsize=10)

xの最近傍の点を探索(d: distance, i: index)

x = (10, 10)
d, i = tree.query(x)
print('most nearest node:', data[i])

xの距離r近傍の点を探索

x = (10, 10); r = 1
res = tree.query_ball_point(x, r)
print('r =', r, 'near nodes', [data[i] for i in res])

距離がr以下のペアを全て探索

r = 1
res = tree.query_pairs(r)
print('r =', r, 'pairs', res)

以上は2次元で考えていますが, より高次元でもkd treeを実行することができます。