最近在讀李航老師的《統(tǒng)計學習方法》，讀到第三章的k近鄰算法時，在N>>k時遍歷搜索比較費時，為了更高效的搜索可以采用kd Tree的方式組織Training數(shù)據(jù)，我看到一篇博客，前面的圖示理解部分說的比較到位，不過博主的代碼有些問題，包括：

• 代碼的完整邏輯是不對的
• 并沒有對已搜索的節(jié)點進行標注，導致重復計算（走回頭路），這樣整個kd Tree的初衷就完全沒意義了

于是在該代碼基礎上改了一版正確的，帶了一些調(diào)試信息。

• 更正代碼邏輯
• 代碼clean up
• 更新為python 3兼容
• 增加up_traced標志避免走回頭路，支持多次搜索，每次搜索前做好該標志的清理

# -*- coding: utf-8 -*-

import numpy as np


class Node:
    def __init__(self, data, parent, dim):
        self.data = data
        self.parent = parent
        self.lChild = None
        self.rChild = None
        self.dim = dim
        # only track search_Up process
        self.up_traced = False

    def setLChild(self, lChild):
        self.lChild = lChild

    def setRChild(self, rChild):
        self.rChild = rChild


class KdTree:
    def __init__(self, train):
        self.root = self.__build(train, 1, None)

    def __build(self, train, depth, parent):  # 遞歸建樹
        (m, k) = train.shape

        if m == 0:
            return None

        train = train[train[:, depth % k].argsort()]

        root = Node(train[m//2], parent, depth % k)
        root.setLChild(self.__build(train[:m//2, :], depth+1, root))
        root.setRChild(self.__build(train[m//2+1:, :], depth+1, root))
        return root

    def findNearestPointAndDistance(self, point):  # 查找與point距離最近的點
        point = np.array(point)
        node = self.__findSmallestSubSpace(point, self.root)
        print("Start node:", node.data)
        return self.__searchUp(point, node, node, np.linalg.norm(point - node.data))

    def __searchUp(self, point, node, nearestPoint, nearestDistance):
        if node.parent is None:
            return [nearestPoint, nearestDistance]

        print("UP:", node.parent.data)
        node.parent.up_traced = True
        distance = np.linalg.norm(node.parent.data - point)
        if distance < nearestDistance:
            nearestDistance = distance
            nearestPoint = node.parent

        distance = np.abs(node.parent.data[node.dim] - point[node.parent.dim])
        if distance < nearestDistance:
            [p, d] = self.__searchDown(point, node.parent)
            if d < nearestDistance:
                nearestDistance = d
                nearestPoint = p

        [p, d] = self.__searchUp(point, node.parent, nearestPoint, nearestDistance)
        if d < nearestDistance:
            nearestDistance = d
            nearestPoint = p

        return [nearestPoint, nearestDistance]

    def __searchDown(self, point, node):

        nearestDistance = np.linalg.norm(node.data - point)
        nearestPoint = node

        print("DOWN:", node.data)
        if node.lChild is not None and node.lChild.up_traced is False:
            [p, d] = self.__searchDown(point, node.lChild)
            if d < nearestDistance:
                nearestDistance = d
                nearestPoint = p

        if node.rChild is not None and node.rChild.up_traced is False:
            [p, d] = self.__searchDown(point, node.rChild)
            if d < nearestDistance:
                nearestDistance = d
                nearestPoint = p

        print("---- ", nearestPoint.data, nearestDistance)
        return [nearestPoint, nearestDistance]

    def __findSmallestSubSpace(self, point, node):  # 找到這個點所在的最小的子空間
        """
        從根節(jié)點出發(fā)，遞歸地向下訪問kd樹。如果point當前維的坐標小于切分點的坐標，則
        移動到左子節(jié)點，否則移動到右子節(jié)點。直到子節(jié)點為葉節(jié)點為止。
        """
        # New search: clean up up_traced flag for all up path nodes
        node.up_traced = False
        if point[node.dim] < node.data[node.dim]:
            if node.lChild is None:
                return node
            else:
                return self.__findSmallestSubSpace(point, node.lChild)
        else:
            if node.rChild is None:
                return node
            else:
                return self.__findSmallestSubSpace(point, node.rChild)


train = np.array([[2, 3], [5, 4], [9, 6], [4, 7], [8, 1], [7, 2]])
train = np.array([[2, 5], [3, 2], [3, 7], [8, 3], [6, 6], [1, 1], [1, 8]])
kdTree = KdTree(train)


target = np.array([6, 4])
print('##### target :', target)
[p, d] = kdTree.findNearestPointAndDistance(target)

print(p.data, d)
print('---------------------')

(m, k) = train.shape
for i in range(m):
    print(train[i], np.linalg.norm(train[i]-target))


target = np.array([2, 2])
print('')
print('##### target :', target)
[p, d] = kdTree.findNearestPointAndDistance(target)

print(p.data, d)
print('---------------------')

for i in range(m):
    print(train[i], np.linalg.norm(train[i]-target))