import copy
class DistIndex:
def __init__(self, distance, index):
self.distance = distance
self.index = index
def __lt__(self, other):
return self.distance < other.distance
class KNNResultSet:
def __init__(self, capacity):
self.capacity = capacity
self.count = 0
self.worst_dist = 1e10
self.dist_index_list = []
for i in range(capacity):
self.dist_index_list.append(DistIndex(self.worst_dist, 0))
self.comparison_counter = 0
def size(self):
return self.count
def full(self):
return self.count == self.capacity
def worstDist(self):
return self.worst_dist
def add_point(self, dist, index):
self.comparison_counter += 1
if dist > self.worst_dist:
return
if self.count < self.capacity:
self.count += 1
i = self.count - 1
while i > 0:
if self.dist_index_list[i-1].distance > dist:
self.dist_index_list[i] = copy.deepcopy(self.dist_index_list[i-1])
i -= 1
else:
break
self.dist_index_list[i].distance = dist
self.dist_index_list[i].index = index
self.worst_dist = self.dist_index_list[self.capacity-1].distance
def __str__(self):
output = ''
for i, dist_index in enumerate(self.dist_index_list):
output += '%d - %.2f\n' % (dist_index.index, dist_index.distance)
output += 'In total %d comparison operations.' % self.comparison_counter
return output
class RadiusNNResultSet:
def __init__(self, radius):
self.radius = radius
self.count = 0
self.worst_dist = radius
self.dist_index_list = []
self.comparison_counter = 0
def size(self):
return self.count
def worstDist(self):
return self.radius
def add_point(self, dist, index):
self.comparison_counter += 1
if dist > self.radius:
return
self.count += 1
self.dist_index_list.append(DistIndex(dist, index))
def __str__(self):
self.dist_index_list.sort()
output = ''
for i, dist_index in enumerate(self.dist_index_list):
output += '%d - %.2f\n' % (dist_index.index, dist_index.distance)
output += 'In total %d neighbors within %f.\nThere are %d comparison operations.' \
% (self.count, self.radius, self.comparison_counter)
return output
import random
import math
import numpy as np
from result_set import KNNResultSet, RadiusNNResultSet
class Node:
def __init__(self, axis, value, left, right, point_indices):
self.axis = axis
self.value = value
self.left = left
self.right = right
self.point_indices = point_indices
def is_leaf(self):
if self.value is None:
return True
else:
return False
def __str__(self):
output = ''
output += 'axis %d, ' % self.axis
if self.value is None:
output += 'split value: leaf, '
else:
output += 'split value: %.2f, ' % self.value
output += 'point_indices: '
output += str(self.point_indices.tolist())
return output
def sort_key_by_vale(key, value):
assert key.shape == value.shape
assert len(key.shape) == 1
sorted_idx = np.argsort(value)
key_sorted = key[sorted_idx]
value_sorted = value[sorted_idx]
return key_sorted, value_sorted
def axis_round_robin(axis, dim):
if axis == dim-1:
return 0
else:
return axis + 1
def kdtree_recursive_build(root, db, point_indices, axis, leaf_size):
if root is None:
root = Node(axis, None, None, None, point_indices)
if len(point_indices) > leaf_size:
point_indices_sorted, _ = sort_key_by_vale(point_indices, db[point_indices, axis])
middle_left_idx=math.ceil(point_indices_sorted.shape[0]/2)-1
middle_left_point_idx=point_indices_sorted[middle_left_idx]
middle_left_point_value=db[middle_left_point_idx,axis]
middle_right_idx=middle_left_idx+1
middle_right_point_idx=point_indices_sorted[middle_right_idx]
middle_right_point_value=db[middle_right_point_idx,axis]
root.value=(middle_left_point_value+middle_right_point_value)*0.5
root.left=kdtree_recursive_build(root.left,db,point_indices_sorted[0:middle_right_idx],axis_round_robin(axis,dim=db.shape[1]),left_size)
root.right=kdtree_recursive_build(root.right,db,point_indices_sorted[middle_right_idx:],axis_round_robin(axis,dim=db.shape[1]),left_size)
return root
def traverse_kdtree(root: Node, depth, max_depth):
depth[0] += 1
if max_depth[0] < depth[0]:
max_depth[0] = depth[0]
if root.is_leaf():
print(root)
else:
traverse_kdtree(root.left, depth, max_depth)
traverse_kdtree(root.right, depth, max_depth)
depth[0] -= 1
def kdtree_construction(db_np, leaf_size):
N, dim = db_np.shape[0], db_np.shape[1]
root = None
root = kdtree_recursive_build(root,
db_np,
np.arange(N),
axis=0,
leaf_size=leaf_size)
return root
def kdtree_knn_search(root: Node, db: np.ndarray, result_set: KNNResultSet, query: np.ndarray):
if root is None:
return False
if root.is_leaf():
leaf_points = db[root.point_indices, :]
diff = np.linalg.norm(np.expand_dims(query, 0) - leaf_points, axis=1)
for i in range(diff.shape[0]):
result_set.add_point(diff[i], root.point_indices[i])
return False
if query[root.axis]<=root.value:
kdtree_knn_search(root.left,db,result_set,query)
if math.fabs(query[root.axis]-root.value)<result_set.worstDist():
kdtree_knn_search(root.right,db,result_set,query)
else:
kdtree_knn_search(root.right,db,result_set,query)
if math.fabs(query[root.axis]-root.value)<result_set.worstDist():
knn_search(root.left,db,result_set,query)
return False
def kdtree_radius_search(root: Node, db: np.ndarray, result_set: RadiusNNResultSet, query: np.ndarray):
if root is None:
return False
if root.is_leaf():
leaf_points = db[root.point_indices, :]
diff = np.linalg.norm(np.expand_dims(query, 0) - leaf_points, axis=1)
for i in range(diff.shape[0]):
result_set.add_point(diff[i], root.point_indices[i])
return False
if query[rooot.axis]<=root.value:
kdtree_radius_search(root.left,db,result_set,query)
if math.fabs(query[root.axis]-root.value)<result_set.worstDist():
kdtree_radius_search(root.right,db,result_set,query)
else:
kdtree_radius_search(root.right,db,result_set,query)
if math.fabs(query[root.axis]-root.value)<result_set.worstDist():
kdtree_radius_search(root.left,db,result_set,query)
return False
def main():
db_size = 64
dim = 3
leaf_size = 4
k = 1
db_np = np.random.rand(db_size, dim)
root = kdtree_construction(db_np, leaf_size=leaf_size)
depth = [0]
max_depth = [0]
traverse_kdtree(root, depth, max_depth)
print("tree max depth: %d" % max_depth[0])
if __name__ == '__main__':
main()