近似最近邻搜索 (二) 树方法

树方法

kd-tree

kd-tree (k dimensional tree )是树方法的经典算法，其是二分搜索树在多维空间的推广。二分搜索树检索迅速的原因是规定将数据中大于当前节点数据的方在一侧(比如右子树)，而不小于的放在另一侧(比如左子树),这样检索数据时，即可获得logn的速度。kd-tree 也类似,也是二叉搜索树，只不过其中的一维数据变成了n维数据。如同x-y轴坐标系将二维空间分成四个部分一样， n维空间也可这样划分。然后，规定将大于分割点的数据放在某一侧(比如右子树)，不小于分割点的数据放在另一侧(比如左子树)。 kd-tree 适用\(N>>2^k\) 的情形，其中N为数据数量。不过实际生产中，k不能太大，一般10维左右时，效果不错，最好不要超过20维。

kd-tree相对于二叉搜索树维度变多了，操作基本一致，只有一点需要注意：因为kd-tree中的数据是多于一维的，那么每次分叉时(即划分)时，需要操作某一维度，因此涉及如何选取维度以及在此维度下用哪个数据点划分 。

为使得数据划分相对均匀些，应选这样的维度：在此维度上，数据非常分散， 分散程度可用方差来表示。方差越大说明数据越分散，也即就数据划分后会相对平衡。(有时也会为了效率根据一定规则选取，比如哈希等)

当选好维度后，可以直选取此维度中位数据划分即可， 有时中位数据不在数据中， 那么只需用与中位数最近数据点即可。

然后迭代上述过程，直至建树完成。

以下简易代码^[1]：

from collections import namedtuple
from operator import itemgetter
from pprint import pformat

class Node(namedtuple("Node", "location left_child right_child")):
    def __repr__(self):
        return pformat(tuple(self))

def kdtree(point_list, depth: int = 0):
    if not point_list:
        return None

    k = len(point_list[0])  # assumes all points have the same dimension
    # Select axis based on depth so that axis cycles through all valid values
    axis = depth % k

    # Sort point list by axis and choose median as pivot element
    point_list.sort(key=itemgetter(axis))
    median = len(point_list) // 2

    # Create node and construct subtrees
    return Node(
        location=point_list[median],
        left_child=kdtree(point_list[:median], depth + 1),
        right_child=kdtree(point_list[median + 1 :], depth + 1),
    )

def main():
    """Example usage"""
    point_list = [(7, 2), (5, 4), (9, 6), (4, 7), (8, 1), (2, 3)]
    tree = kdtree(point_list)
    print(tree)

if __name__ == "__main__":
    main()

了解算法后，使用的话，可以根据需要自己写，另一个是用一些开源的实现比较好的,比如sklearn^[2]：

!/usr/bin/python
# -*- coding: UTF-8 -*-
import numpy as np
from matplotlib import pyplot as plt
from matplotlib.patches import Circle
from sklearn.neighbors import KDTree

np.random.seed(0)
points = np.random.random((100, 2))
tree = KDTree(points)
point = points[0]

# kNN
dists, indices = tree.query([point], k=3)
print(dists, indices)

# query radius
indices = tree.query_radius([point], r=0.2)
print(indices)

fig = plt.figure()
ax = fig.add_subplot(111, aspect='equal')
ax.add_patch(Circle(point, 0.2, color='r', fill=False))
X, Y = [p[0] for p in points], [p[1] for p in points]
plt.scatter(X, Y)
plt.scatter([point[0]], [point[1]], c='r')
plt.show()

---

1. https://en.wikipedia.org/wiki/K-d_tree ↩︎

2. https://leileiluoluo.com/posts/kdtree-algorithm-and-implementation.html ↩︎