kdtreepython_Python实现KDTree最近邻算法

这里将写了一个KDTree类&＃xff0c;仅实现了最近邻&＃xff0c;K近邻之后若有时间再更新&＃xff1a;

from collections import namedtuple

from operator import itemgetter

from pprint import pformat

import numpy as np

class Node(namedtuple(&＃39;Node&＃39;, &＃39;location left_child right_child&＃39;)):

def __repr__(self):

return pformat(tuple(self))

class KDTree():

def __init__(self, points):

self.tree &＃61; self._make_kdtree(points)

if len(points) > 0:

self.k &＃61; len(points[0])

else:

self.k &＃61; None

def _make_kdtree(self, points, depth&＃61;0):

if not points:

return None

k &＃61; len(points[0])

axis &＃61; depth % k

points.sort(key&＃61;itemgetter(axis))

median &＃61; len(points) // 2

return Node(

location&＃61;points[median],

left_child&＃61;self._make_kdtree(points[:median], depth &＃43; 1),

right_child&＃61;self._make_kdtree(points[median &＃43; 1:], depth &＃43; 1))

def find_nearest(self,

point,

root&＃61;None,

axis&＃61;0,

dist_func&＃61;lambda x, y: np.linalg.norm(x - y)):

if root is None:

root &＃61; self.tree

self._best &＃61; None

# 若不是叶节点&＃xff0c;则继续向下走

if root.left_child or root.right_child:

new_axis &＃61; (axis &＃43; 1) % self.k

if point[axis]

self.find_nearest(point, root.left_child, new_axis)

elif root.right_child:

self.find_nearest(point, root.right_child, new_axis)

# 回溯&＃xff1a;尝试更新 best

dist &＃61; dist_func(root.location, point)

if self._best is None or dist

self._best &＃61; (dist, root.location)

# 若超球与另一边超矩形相交

if abs(point[axis] - root.location[axis])

new_axis &＃61; (axis &＃43; 1) % self.k

if root.left_child and point[axis] >&＃61; root.location[axis]:

self.find_nearest(point, root.left_child, new_axis)

elif root.right_child and point[axis]

self.find_nearest(point, root.right_child, new_axis)

return self._best

测试&＃xff1a;

point_list &＃61; [(2, 3, 3), (5, 4, 4), (9, 6, 7), (4, 7, 7), (8, 1, 1), (7, 2, 2)]

kdtree &＃61; KDTree(point_list)

point &＃61; np.array([5, 5, 5])

print(kdtree.find_nearest(point))

输出&＃xff1a;

(1.4142135623730951, (5, 4, 4))

与 Scikit-Learn 性能对比(上是我的实现&＃xff0c;下是 Scikit-Learn 的实现)&＃xff1a;

可以看到仅相差 1 毫秒&＃xff0c;所以性能说得过去。

(本文完)