1 import numpy as np
2 arr = np.array([[2, 3], [5, 4], [9, 6], [4, 7], [8, 1], [7, 2]])
3 arr.shape
4
5 class KDTree():
6 def __init__(self):
7 self.value = None
8 self.left = None
9 self.right = None
10 self.axis = None
11
12 def create(arr, k, h=0):
13 if arr.shape[0] == 0:
14 return None
15 tree = KDTree()
16 axis = h % k
17
18 if arr.shape[0] == 1:
19 tree.value = arr[0]
20 tree.left = None
21 tree.right = None
22 tree.axis = axis
23 else:
24 arr = sorted(arr, key = lambda x:x[axis])
25 arr = np.array(arr)
26 i = arr.shape[0]//2
27
28 tree.value = arr[i]
29 tree.left = create(arr[0:i], k, h+1)
30 tree.right = create(arr[i+1:], k, h+1)
31 tree.axis = axis
32 return tree
33
34 k = KDTree()
35
36 k = create(arr, arr.shape[1])
37
38 def preOrder(k):
39 print('当前节点:' + str(k.value))
40
41 if k.left:
42 preOrder(k.left)
43 if k.right:
44 preOrder(k.right)
45
46 preOrder(k)
47
48 def dis(a, b):
49 return np.linalg.norm(a-b)
50 def search(kd, goal, k, h=0):
51 '''输入：kd树，目标点、特征维度k以及当前深度h'''
52 '''输出：在kd树上的与目标点距离（欧氏距离）最近的距离'''
53 if kd.left == None and kd.right == None:
54 return dis(goal, kd.value)
55 if kd.left == None:
56 return min(search(kd.right, goal, k, h+1), dis(kd.value, goal))
57 if kd.right == None:
58 return min(search(kd.left, goal, k, h+1), dis(kd.value, goal))
59 axis = h%k
60
61 if goal[axis] < kd.value[axis]:
62 cur_dis = search(kd.left, goal, k, h+1)
63 else:
64 cur_dis = search(kd.right, goal, k, h+1)
65
66
67 if cur_dis < kd.value[axis]-goal[axis]:////cut 取绝对值
68 return cur_dis;
69 else:
70 if goal[axis] < kd.value[axis]:
71 cur_dis = min(search(kd.right, goal, k, h+1), cur_dis, dis(kd.value, goal))
72 else:
73 cur_dis = min(search(kd.left, goal, k, h+1), cur_dis, dis(kd.value, goal))
74 return cur_dis
75
76 search(k, np.array([9, 6]), 2)
