Implemented KD-Tree Data Structure

Author: Ramy-Badr-Ahmed

data_structures/kd_tree/__init__.py

@@ -0,0 +1,19 @@
+from .kd_node import KDNode
+
+def build_kdtree(points, depth=0):
+    if not points:
+        return None
+
+    k = len(points[0])  # dimensionality of the points
+    axis = depth % k
+
+    # Sort point list and choose median as pivot element
+    points.sort(key=lambda x: x[axis])
+    median_idx = len(points) // 2
+
+    # Create node and construct subtrees
+    return KDNode(
+        point = points[median_idx],
+        left = build_kdtree(points[:median_idx], depth + 1),
+        right = build_kdtree(points[median_idx + 1:], depth + 1)
+    )

data_structures/kd_tree/build_kdtree.py

@@ -0,0 +1,22 @@
+import numpy as np
+
+from hypercube_points import hypercube_points
+from data_structures.kd_tree.build_kdtree import build_kdtree
+from data_structures.kd_tree.nearest_neighbour_search import nearest_neighbour_search
+
+
+num_points = 5000
+cube_size = 10
+num_dimensions = 10
+
+points = hypercube_points(num_points, cube_size, num_dimensions)
+hypercube_kdtree = build_kdtree(points.tolist())
+
+query_point = np.random.rand(num_dimensions).tolist()
+
+nearest_point, nearest_dist, nodes_visited = nearest_neighbour_search(hypercube_kdtree, query_point)
+
+print(f"Query point: {query_point}")
+print(f"Nearest point: {nearest_point}")
+print(f"Distance: {nearest_dist:.4f}")
+print(f"Nodes visited: {nodes_visited}")

data_structures/kd_tree/example/example_usage.py

@@ -0,0 +1,4 @@
+import numpy as np
+
+def hypercube_points(num_points, hypercube_size, num_dimensions):
+    return hypercube_size * np.random.rand(num_points, num_dimensions)

data_structures/kd_tree/example/hypercube_points.py

@@ -0,0 +1,5 @@
+class KDNode:
+    def __init__(self, point, left = None, right = None):
+        self.point = point
+        self.left = left
+        self.right = right

data_structures/kd_tree/kd_node.py

@@ -0,0 +1,41 @@
+def nearest_neighbour_search(root, query_point):
+    nearest_point = None
+    nearest_dist = float('inf')
+    nodes_visited = 0
+
+    def search(node, depth=0):
+        nonlocal nearest_point, nearest_dist, nodes_visited
+        if node is None:
+            return
+
+        nodes_visited += 1
+
+        # Calculate the current distance (squared distance)
+        current_point = node.point
+        current_dist = sum((qp - cp) ** 2 for qp, cp in zip(query_point, current_point))
+
+        # Update nearest point if the current node is closer
+        if nearest_point is None or current_dist < nearest_dist:
+            nearest_point = current_point
+            nearest_dist = current_dist
+
+        # Determine which subtree to search first (based on axis and query point)
+        k = len(query_point)  # dimensionality of points
+        axis = depth % k
+
+        if query_point[axis] <= current_point[axis]:
+            nearer_subtree = node.left
+            further_subtree = node.right
+        else:
+            nearer_subtree = node.right
+            further_subtree = node.left
+
+        # Search the nearer subtree first
+        search(nearer_subtree, depth + 1)
+
+        # If the further subtree has a closer point
+        if (query_point[axis] - current_point[axis]) ** 2 < nearest_dist:
+            search(further_subtree, depth + 1)
+
+    search(root, 0)
+    return nearest_point, nearest_dist, nodes_visited