← retour aux snippets

scipy.spatial: KDTree voisins

Trouver les plus proches voisins avec KDTree.

python stats #scipy#kdtree#nearest

objectif

Trouver les plus proches voisins avec KDTree.

code minimal

import numpy as np
from scipy.spatial import KDTree

pts = np.random.default_rng(0).normal(size=(10,2))
tree = KDTree(pts)
dist, idx = tree.query([0.0,0.0], k=3)
print(int(idx[0]))

utilisation

print(tree.query_ball_point([0.0,0.0], r=1.0)[:3])

variante(s) utile(s)

import numpy as np
q = np.array([[0.1,0.2],[1.0,1.0]])
print(tree.query(q, k=1)[1].tolist())

notes

  • KDTree rapide en faible dimension; sinon BallTree.