docstring for search()

Author: Ramy-Badr-Ahmed

data_structures/kd_tree/nearest_neighbour_search.py

@@ -1,7 +1,10 @@
 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]:
+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.
 
@@ -20,6 +23,18 @@ def nearest_neighbour_search(root: Optional[KDNode], query_point: List[float]) -
     nodes_visited: int = 0
 
     def search(node: Optional[KDNode], depth: int = 0) -> None:
+        """
+        Recursively searches the KD-Tree to find the nearest point to the query point.
+
+        Args:
+            node (Optional[KDNode]): The current node being examined.
+            depth (int): The current depth of the tree, which determines the axis to split on.
+
+        Updates:
+            nearest_point: The closest point found so far in the KD-Tree.
+            nearest_dist: The squared distance from the query point to the nearest point found.
+            nodes_visited: The number of nodes visited during the search.
+        """
         nonlocal nearest_point, nearest_dist, nodes_visited
         if node is None:
             return
@@ -28,7 +43,9 @@ def search(node: Optional[KDNode], depth: int = 0) -> None:
 
         # Calculate the current distance (squared distance)
         current_point = node.point
-        current_dist = sum((query_coord - point_coord) ** 2 for query_coord, point_coord 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:
@@ -49,7 +66,7 @@ def search(node: Optional[KDNode], depth: int = 0) -> None:
         # Search the nearer subtree first
         search(nearer_subtree, depth + 1)
 
-        # If the further subtree has a closer point
+        # If the further subtree has a closer point, search it
         if (query_point[axis] - current_point[axis]) ** 2 < nearest_dist:
             search(further_subtree, depth + 1)