mirror of
https://github.com/AtsushiSakai/PythonRobotics.git
synced 2026-04-22 03:00:22 -04:00
50 lines
943 B
Python
50 lines
943 B
Python
"""
|
|
|
|
Kd tree Search library
|
|
|
|
author: Atsushi Sakai (@Atsushi_twi)
|
|
|
|
"""
|
|
|
|
import scipy.spatial
|
|
|
|
|
|
class KDTree:
|
|
"""
|
|
Nearest neighbor search class with KDTree
|
|
"""
|
|
|
|
def __init__(self, data):
|
|
# store kd-tree
|
|
self.tree = scipy.spatial.cKDTree(data)
|
|
|
|
def search(self, inp, k=1):
|
|
"""
|
|
Search NN
|
|
|
|
inp: input data, single frame or multi frame
|
|
|
|
"""
|
|
|
|
if len(inp.shape) >= 2: # multi input
|
|
index = []
|
|
dist = []
|
|
|
|
for i in inp.T:
|
|
idist, iindex = self.tree.query(i, k=k)
|
|
index.append(iindex)
|
|
dist.append(idist)
|
|
|
|
return index, dist
|
|
|
|
dist, index = self.tree.query(inp, k=k)
|
|
return index, dist
|
|
|
|
def search_in_distance(self, inp, r):
|
|
"""
|
|
find points with in a distance r
|
|
"""
|
|
|
|
index = self.tree.query_ball_point(inp, r)
|
|
return index
|