kd-tree找最邻近点 Python实现

kd-tree找最邻近点 Python实现

基本概念

kd-tree是KNN算法的一种实现。算法的基本思想是用多维空间中的实例点，将空间划分为多块，成二叉树形结构。划分超矩形上的实例点是树的非叶子节点，而每个超矩形内部的实例点是叶子结点。

超矩形划分方法

有数据集datalist，其中的数据是Xi,每个Xi由多个特征值组成。首先将所有数据的Xi[0]找出，取得Xi[0]的中位数center，在树的根节点中保存Xi[0] == center的实例点Xi,树根的左子树递归构造Xi[0] < center的数据集，树根的右子树构造X[0] > center的数据集。同时在第二层，划分特征变为Xi[1], 划分特征随着树的深度改变，为 (d - 1) % k , k是特征的维度，d是此时划分的树深度。

python实现

import numpy as np
import matplotlib.pyplot as plot
import math
# kdtree类


class KdTree(object):
    """docstring for KdTree."""

    def __str__(self):
        return '{ nodes:' + str(self.nodes) + ', left:' + str(self.l) + ', right:' + str(self.r) + '}'

    def __init__(self):
        self.split = None
        self.l = None
        self.r = None
        self.f = None
        self.nodes = []


def createKdTree(split, datalist, k):
    if datalist is None or len(datalist) == 0:
        return None
    split = split % k
    node = KdTree()
    # 求中位数
    center = np.sort([a[split] for a in datalist])[int(len(datalist)/2)]
    leftData = [a for a in datalist if a[split] < center]
    rightData