K-D Tree(K-Dimensional Tree)是一种用于处理 k 维空间数据的数据结构。它可以用于搜索最近邻点、范围搜索等问题。下面是一个简单的 Python 实现。
首先,我们需要一个节点类来表示树中的节点:
```python
class Node:
def __init__(self, point=None, axis=None, left=None, right=None):
self.point = point # 节点代表的点
self.axis = axis # 划分轴
self.left = left # 左子树
self.right = right # 右子树
```
接下来,我们需要一个 `build_kdtree` 函数来构建 K-D Tree:
```python
def build_kdtree(points, depth=0):
n = len(points)
if n <= 0:
return None
axis = depth % k # 选择轴
points.sort(key=lambda point: point[axis])
mid = n // 2
return Node(
point=points[mid],
axis=axis,
left=build_kdtree(points[:mid], depth+1),
right=build_kdtree(points[mid+1:], depth+1)
)
```
这个函数会递归地将点集分成两部分,构建 K-D Tree。
接下来,我们需要实现两个重要的子函数:`nearest_neighbor` 和 `range_search`。
`nearest_neighbor` 函数用于查找 K-D Tree 中最近的点:
```python
def nearest_neighbor(root, point):
best = None
best_distance = float("inf")
def search(node, depth=0):
nonlocal best, best_distance
if node is None:
return
distance = euclidean_distance(point, node.point)
if distance < best_distance:
best = node.point
best_distance = distance
axis = depth % k
if point[axis] < node.point[axis]:
search(node.left, depth+1)
if point[axis] + best_distance >= node.point[axis]:
search(node.right, depth+1)
else:
search(node.right, depth+1)
if point[axis] - best_distance <= node.point[axis]:
search(node.left, depth+1)
search(root)
return best
```
这个函数会从根节点开始递归地搜索 K-D Tree,找到与目标点最近的点。
最后,我们需要实现 `range_search` 函数来查找 K-D Tree 中给定范围内的所有点:
```python
def range_search(root, lower, upper):
result = []
def search(node, depth=0):
if node is None:
return
if lower <= node.point <= upper:
result.append(node.point)
axis = depth % k
if lower[axis] <= node.point[axis]:
search(node.left, depth+1)
if upper[axis] >= node.point[axis]:
search(node.right, depth+1)
search(root)
return result
```
这个函数会从根节点开始递归地搜索 K-D Tree,找到所有满足条件的点。
完整的代码如下所示:
```python
import math
k = 2
class Node:
def __init__(self, point=None, axis=None, left=None, right=None):
self.point = point # 节点代表的点
self.axis = axis # 划分轴
self.left = left # 左子树
self.right = right # 右子树
def euclidean_distance(p, q):
return math.sqrt(sum((p[i] - q[i])**2 for i in range(k)))
def build_kdtree(points, depth=0):
n = len(points)
if n <= 0:
return None
axis = depth % k # 选择轴
points.sort(key=lambda point: point[axis])
mid = n // 2
return Node(
point=points[mid],
axis=axis,
left=build_kdtree(points[:mid], depth+1),
right=build_kdtree(points[mid+1:], depth+1)
)
def nearest_neighbor(root, point):
best = None
best_distance = float("inf")
def search(node, depth=0):
nonlocal best, best_distance
if node is None:
return
distance = euclidean_distance(point, node.point)
if distance < best_distance:
best = node.point
best_distance = distance
axis = depth % k
if point[axis] < node.point[axis]:
search(node.left, depth+1)
if point[axis] + best_distance >= node.point[axis]:
search(node.right, depth+1)
else:
search(node.right, depth+1)
if point[axis] - best_distance <= node.point[axis]:
search(node.left, depth+1)
search(root)
return best
def range_search(root, lower, upper):
result = []
def search(node, depth=0):
if node is None:
return
if lower <= node.point <= upper:
result.append(node.point)
axis = depth % k
if lower[axis] <= node.point[axis]:
search(node.left, depth+1)
if upper[axis] >= node.point[axis]:
search(node.right, depth+1)
search(root)
return result
```
这是一个简单的 K-D Tree 的 Python 实现,可以用来解决搜索最近邻点、范围搜索等问题。