def test_find_knn_from_kd_tree():
k = 3
samples = np.array([[2, 3], [5, 4], [9, 6], [4, 7], [8, 1], [7, 2]])
ls = np.array(list(range(6)))
kd_tree = KDTree.construct_kd_tree(samples, ls)
sample = np.array([4, 6])
knn_heap = KDTree.find_knn_from_kd_tree(sample, kd_tree, k)
print(knn_heap.knns)
def test_search_sub_tree():
capacity = 2
samples = np.array([[2, 3], [5, 4], [9, 6], [4, 7], [8, 1], [7, 2]])
ls = np.array(list(range(6)))
kd_tree = KDTree.construct_kd_tree(samples, ls)
knn_heap = KDTree.KnnHeap(capacity)
KDTree.search_sub_tree((0, 0), kd_tree, knn_heap)
print(knn_heap.knns)
def get_kd_tree(train_samples):
global kd_tree
if kd_tree is None:
# print(train_samples.dtype)
m, _ = train_samples.shape
ls = np.array(list(range(m)))
kd_tree = KDTree.construct_kd_tree(train_samples, ls)
return kd_tree
def test_construct_kd_tree():
samples = np.array([[2, 3], [5, 4], [9, 6], [4, 7], [8, 1], [7, 2]])
ls = np.array(list(range(6)))
kd_tree = KDTree.construct_kd_tree(samples, ls)
print_kd_tree(kd_tree)
