added typehints and docstrings

Author: Ramy-Badr-Ahmed

data_structures/kd_tree/build_kdtree.py

@@ -1,20 +1,30 @@
+from typing import List, Optional
 from .kd_node import KDNode
 
+def build_kdtree(points: List[List[float]], depth: int = 0) -> Optional[KDNode]:
+    """
+    Builds a KD-Tree from a set of k-dimensional points.
 
-def build_kdtree(points, depth=0):
+    Args:
+        points (List[List[float]]): A list of k-dimensional points (each point is a list of floats).
+        depth (int): The current depth in the tree. Used to determine the splitting axis. Defaults to 0.
+
+    Returns:
+        Optional[KDNode]: The root of the KD-Tree or None if the input list is empty.
+    """
     if not points:
         return None
 
-    k = len(points[0])  # dimensionality of the points
+    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])
+    points.sort(key=lambda point: point[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),
+        right=build_kdtree(points[median_idx + 1:], depth + 1),
     )

data_structures/kd_tree/example/example_usage.py

@@ -1,25 +1,36 @@
 import numpy as np
-
+from typing import List
 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
 
+def main() -> None:
+    """
+    Demonstrates the use of KD-Tree by building it from random points
+    in a 10-dimensional hypercube and performing a nearest neighbor search.
+    """
+    num_points: int = 5000
+    cube_size: int = 10
+    num_dimensions: int = 10
 
-num_points = 5000
-cube_size = 10
-num_dimensions = 10
+    # Generate random points within the hypercube
+    points: np.ndarray = hypercube_points(num_points, cube_size, num_dimensions)
+    hypercube_kdtree = build_kdtree(points.tolist())
 
-points = hypercube_points(num_points, cube_size, num_dimensions)
-hypercube_kdtree = build_kdtree(points.tolist())
+    # Generate a random query point within the same space
+    rng = np.random.default_rng()
+    query_point: List[float] = rng.random(num_dimensions).tolist()
 
-rng = np.random.default_rng()
-query_point = rng.random(num_dimensions).tolist()
+    # Perform nearest neighbor search
+    nearest_point, nearest_dist, nodes_visited = nearest_neighbour_search(
+        hypercube_kdtree, query_point
+    )
 
-nearest_point, nearest_dist, nodes_visited = nearest_neighbour_search(
-    hypercube_kdtree, query_point
-)
+    # Print the results
+    print(f"Query point: {query_point}")
+    print(f"Nearest point: {nearest_point}")
+    print(f"Distance: {nearest_dist:.4f}")
+    print(f"Nodes visited: {nodes_visited}")
 
-print(f"Query point: {query_point}")
-print(f"Nearest point: {nearest_point}")
-print(f"Distance: {nearest_dist:.4f}")
-print(f"Nodes visited: {nodes_visited}")
+if __name__ == "__main__":
+    main()

data_structures/kd_tree/example/hypercube_points.py

@@ -1,6 +1,17 @@
 import numpy as np
+from typing import Union
 
+def hypercube_points(num_points: int, hypercube_size: Union[int, float], num_dimensions: int) -> np.ndarray:
+    """
+    Generates random points uniformly distributed within an n-dimensional hypercube.
 
-def hypercube_points(num_points, hypercube_size, num_dimensions):
+    Args:
+        num_points (int): The number of random points to generate.
+        hypercube_size (Union[int, float]): The size of the hypercube (side length).
+        num_dimensions (int): The number of dimensions of the hypercube.
+
+    Returns:
+        np.ndarray: An array of shape (num_points, num_dimensions) with the generated points.
+    """
     rng = np.random.default_rng()
     return hypercube_size * rng.random((num_points, num_dimensions))

data_structures/kd_tree/kd_node.py

@@ -1,5 +1,24 @@
+from typing import List, Optional
+
 class KDNode:
-    def __init__(self, point, left=None, right=None):
+    """
+    Represents a node in a KD-Tree.
+
+    Attributes:
+        point (List[float]): The k-dimensional point stored in this node.
+        left (Optional[KDNode]): The left subtree of this node.
+        right (Optional[KDNode]): The right subtree of this node.
+    """
+
+    def __init__(self, point: List[float], left: Optional['KDNode'] = None, right: Optional['KDNode'] = None) -> None:
+        """
+        Initializes a KDNode with a point and optional left and right children.
+
+        Args:
+            point (List[float]): The k-dimensional point to be stored in this node.
+            left (Optional[KDNode]): The left subtree of this node. Defaults to None.
+            right (Optional[KDNode]): The right subtree of this node. Defaults to None.
+        """
         self.point = point
         self.left = left
         self.right = right

data_structures/kd_tree/nearest_neighbour_search.py

@@ -1,9 +1,25 @@
-def nearest_neighbour_search(root, query_point):
-    nearest_point = None
-    nearest_dist = float("inf")
-    nodes_visited = 0
-
-    def search(node, depth=0):
+from typing import Optional, List, Tuple
+from .kd_node import KDNode
+
+def nearest_neighbour_search(root: Optional[KDNode], query_point: List[float]) -> Tuple[Optional[List[float]], float, int]:
+    """
+    Performs a nearest neighbor search in a KD-Tree for a given query point.
+
+    Args:
+        root (Optional[KDNode]): The root node of the KD-Tree.
+        query_point (List[float]): The point for which the nearest neighbor is being searched.
+
+    Returns:
+        Tuple[Optional[List[float]], float, int]:
+            - The nearest point found in the KD-Tree to the query point.
+            - The squared distance to the nearest point.
+            - The number of nodes visited during the search.
+    """
+    nearest_point: Optional[List[float]] = None
+    nearest_dist: float = float("inf")
+    nodes_visited: int = 0
+
+    def search(node: Optional[KDNode], depth: int = 0) -> None:
         nonlocal nearest_point, nearest_dist, nodes_visited
         if node is None:
             return
@@ -12,7 +28,7 @@ def search(node, depth=0):
 
         # 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))
+        current_dist = sum((query_coord - point_coord) ** 2 for query_coord, point_coord 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: