updated tests and used | for type annotations

Author: Ramy-Badr-Ahmed

data_structures/kd_tree/build_kdtree.py

@@ -1,8 +1,8 @@
-from typing import Optional
-from .kd_node import KDNode
+from data_structures.kd_tree.kd_node import KDNode
 
-
-def build_kdtree(points: list[list[float]], depth: int = 0) -> Optional[KDNode]:
+def build_kdtree(
+        points: list[list[float]], depth: int = 0
+) -> KDNode | None:
     """
     Builds a KD-Tree from a list of points.
 
@@ -11,7 +11,7 @@ def build_kdtree(points: list[list[float]], depth: int = 0) -> Optional[KDNode]:
         depth (int): The current depth in the tree (used to determine axis for splitting).
 
     Returns:
-        Optional[KDNode]: The root node of the KD-Tree.
+        KDNode | None: The root node of the KD-Tree, or None if no points are provided.
     """
     if not points:
         return None
@@ -24,8 +24,11 @@ def build_kdtree(points: list[list[float]], depth: int = 0) -> Optional[KDNode]:
     median_idx = len(points) // 2
 
     # Create node and construct subtrees
+    left_points = points[:median_idx]
+    right_points = points[median_idx + 1:]
+
     return KDNode(
         point=points[median_idx],
-        left=build_kdtree(points[:median_idx], depth + 1),
-        right=build_kdtree(points[median_idx + 1 :], depth + 1),
+        left=build_kdtree(left_points, depth + 1),
+        right=build_kdtree(right_points, depth + 1),
     )

data_structures/kd_tree/example/example_usage.py

@@ -1,5 +1,5 @@
 import numpy as np
-from hypercube_points import hypercube_points
+from data_structures.kd_tree.example.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
 

data_structures/kd_tree/example/hypercube_points.py

@@ -1,8 +1,7 @@
 import numpy as np
 
-
 def hypercube_points(
-    num_points: int, hypercube_size: float, num_dimensions: int
+        num_points: int, hypercube_size: float, num_dimensions: int
 ) -> np.ndarray:
     """
     Generates random points uniformly distributed within an n-dimensional hypercube.
@@ -16,4 +15,5 @@ def hypercube_points(
         np.ndarray: An array of shape (num_points, num_dimensions) with generated points.
     """
     rng = np.random.default_rng()
-    return hypercube_size * rng.random((num_points, num_dimensions))
+    shape = (num_points, num_dimensions)
+    return hypercube_size * rng.random(shape)

data_structures/kd_tree/kd_node.py

@@ -1,6 +1,5 @@
 from typing import Optional
 
-
 class KDNode:
     """
     Represents a node in a KD-Tree.

data_structures/kd_tree/nearest_neighbour_search.py

@@ -1,34 +1,36 @@
-from typing import Optional
 from data_structures.kd_tree.kd_node import KDNode
 
-
 def nearest_neighbour_search(
-    root: Optional[KDNode], query_point: list[float]
-) -> tuple[Optional[list[float]], float, int]:
+        root: KDNode | None,
+        query_point: list[float]
+) -> tuple[list[float] | None, 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.
+        root (KDNode | None): 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.
+        tuple[list[float] | None, float, int]:
+            - The nearest point found in the KD-Tree to the query point, or None if no point is found.
             - The squared distance to the nearest point.
             - The number of nodes visited during the search.
     """
-    nearest_point: Optional[list[float]] = None
+    nearest_point: list[float] | None = None
     nearest_dist: float = float("inf")
     nodes_visited: int = 0
 
-    def search(node: Optional[KDNode], depth: int = 0) -> None:
+    def search(
+            node: KDNode | None,
+            depth: int = 0
+    ) -> None:
         """
-        Recursively searches the KD-Tree for the nearest neighbor.
+        Recursively searches for the nearest neighbor in the KD-Tree.
 
         Args:
-            node (Optional[KDNode]): The current node in the KD-Tree.
-            depth (int): The current depth in the tree.
+            node (KDNode | None): The current node in the KD-Tree.
+            depth (int): The current depth in the KD-Tree.
         """
         nonlocal nearest_point, nearest_dist, nodes_visited
         if node is None:

data_structures/kd_tree/tests/test_kdtree.py

@@ -1,73 +1,70 @@
-import unittest
 import numpy as np
 from data_structures.kd_tree.build_kdtree import build_kdtree
 from data_structures.kd_tree.nearest_neighbour_search import nearest_neighbour_search
 from data_structures.kd_tree.kd_node import KDNode
 from data_structures.kd_tree.example.hypercube_points import hypercube_points
 
+def test_build_kdtree():
+    """
+    Test that KD-Tree is built correctly.
+    """
+    num_points = 10
+    cube_size = 10.0
+    num_dimensions = 2
+    points = hypercube_points(num_points, cube_size, num_dimensions)
+    kdtree = build_kdtree(points.tolist())
 
-class TestKDTree(unittest.TestCase):
-    def setUp(self):
-        """
-        Set up test data.
-        """
-        self.num_points = 10
-        self.cube_size = 10.0
-        self.num_dimensions = 2
-        self.points = hypercube_points(
-            self.num_points, self.cube_size, self.num_dimensions
-        )
-        self.kdtree = build_kdtree(self.points.tolist())
+    # Check if root is not None
+    assert kdtree is not None
 
-    def test_build_kdtree(self):
-        """
-        Test that KD-Tree is built correctly.
-        """
-        # Check if root is not None
-        self.assertIsNotNone(self.kdtree)
+    # Check if root has correct dimensions
+    assert len(kdtree.point) == num_dimensions
 
-        # Check if root has correct dimensions
-        self.assertEqual(len(self.kdtree.point), self.num_dimensions)
+    # Check that the tree is balanced to some extent (simplistic check)
+    assert isinstance(kdtree, KDNode)
 
-        # Check that the tree is balanced to some extent (simplistic check)
-        self.assertIsInstance(self.kdtree, KDNode)
+def test_nearest_neighbour_search():
+    """
+    Test the nearest neighbor search function.
+    """
+    num_points = 10
+    cube_size = 10.0
+    num_dimensions = 2
+    points = hypercube_points(num_points, cube_size, num_dimensions)
+    kdtree = build_kdtree(points.tolist())
 
-    def test_nearest_neighbour_search(self):
-        """
-        Test the nearest neighbor search function.
-        """
-        rng = np.random.default_rng()
-        query_point = rng.random(self.num_dimensions).tolist()
+    rng = np.random.default_rng()
+    query_point = rng.random(num_dimensions).tolist()
 
-        nearest_point, nearest_dist, nodes_visited = nearest_neighbour_search(
-            self.kdtree, query_point
-        )
+    nearest_point, nearest_dist, nodes_visited = nearest_neighbour_search(
+        kdtree, query_point
+    )
 
-        # Check that nearest point is not None
-        self.assertIsNotNone(nearest_point)
+    # Check that nearest point is not None
+    assert nearest_point is not None
 
-        # Check that distance is a non-negative number
-        self.assertGreaterEqual(nearest_dist, 0)
+    # Check that distance is a non-negative number
+    assert nearest_dist >= 0
 
-        # Check that nodes visited is a non-negative integer
-        self.assertGreaterEqual(nodes_visited, 0)
+    # Check that nodes visited is a non-negative integer
+    assert nodes_visited >= 0
 
-    def test_edge_cases(self):
-        """
-        Test edge cases such as an empty KD-Tree.
-        """
-        empty_kdtree = build_kdtree([])
-        query_point = [0.0] * self.num_dimensions
+def test_edge_cases():
+    """
+    Test edge cases such as an empty KD-Tree.
+    """
+    empty_kdtree = build_kdtree([])
+    query_point = [0.0] * 2  # Using a default 2D query point
 
-        nearest_point, nearest_dist, nodes_visited = nearest_neighbour_search(
-            empty_kdtree, query_point
-        )
-
-        # With an empty KD-Tree, nearest_point should be None
-        self.assertIsNone(nearest_point)
-        self.assertEqual(nearest_dist, float("inf"))
-        self.assertEqual(nodes_visited, 0)
+    nearest_point, nearest_dist, nodes_visited = nearest_neighbour_search(
+        empty_kdtree, query_point
+    )
 
+    # With an empty KD-Tree, nearest_point should be None
+    assert nearest_point is None
+    assert nearest_dist == float("inf")
+    assert nodes_visited == 0
 
 if __name__ == "__main__":
-    unittest.main()
+    import pytest
+    pytest.main()