geokdtree.core

  1def kdtree(points, depth, axis_count):
  2    """
  3    Function:
  4
  5    - Build a KDTree from a list of points.
  6
  7    Required Arguments:
  8
  9    - `points`
 10        - Type: list of tuples
 11        - What: A list of points to build the KDTree from
 12
 13    Optional Arguments:
 14
 15    - `depth`
 16        - Type: int
 17        - What: The current depth in the tree (used for axis selection)
 18    - `axis_count`
 19        - Type: int
 20        - What: The number of dimensions in the points (default is 2 for 2D points)
 21
 22    Returns:
 23
 24    - The constructed KDTree as a tuple in the format (point, axis, left, right).
 25    - Where left and right are subtrees.
 26    """
 27    if not points:
 28        return 0
 29    axis = depth % axis_count
 30    points.sort(key=lambda p: p[axis])
 31    median = len(points) // 2
 32    return (
 33        points[median],
 34        axis,
 35        kdtree(points=points[:median], depth=depth + 1, axis_count=axis_count),
 36        kdtree(
 37            points=points[median + 1 :], depth=depth + 1, axis_count=axis_count
 38        ),
 39    )
 40
 41
 42def squared_distance(p1, p2, axis_count=2):
 43    """
 44    Function:
 45
 46    - Calculate the squared distance between two points.
 47
 48    Required Arguments:
 49
 50    - `p1`
 51        - Type: tuple
 52        - What: The first point
 53    - `p2`
 54        - Type: tuple
 55        - What: The second point
 56
 57    Optional Arguments:
 58
 59    - `axis_count`
 60        - Type: int
 61        - What: The number of dimensions in the points
 62        - Default: 2 (for 2D points)
 63
 64    Returns:
 65
 66    - The squared distance between the two points.
 67    """
 68    return sum([(p1[i] - p2[i]) ** 2 for i in range(axis_count)])
 69
 70
 71def closest_point(node, point, best=None, axis_count=2, best_dist=float("inf")):
 72    """
 73    Function:
 74
 75    - Find the closest point in the KDTree to a given point.
 76
 77    Required Arguments:
 78
 79    - `node`
 80        - Type: tuple
 81        - What: The node of the KDTree
 82    - `point`
 83        - Type: tuple
 84        - What: The point to find the closest point to
 85
 86    Optional Arguments:
 87
 88    - `best`
 89        - Type: tuple or None
 90        - What: The best point found so far (default is None)
 91    - `axis_count`
 92        - Type: int
 93        - What: The number of dimensions in the points (default is 2 for 2D points)
 94    - `best_dist`
 95        - Type: float
 96        - What: The best distance found so far (default is infinity)
 97
 98    Returns:
 99
100    - The closest point found in the KDTree to the given point.
101    """
102    if node == 0:
103        return best, best_dist
104    # Get the median node and its distance
105    median_node = node[0]
106    median_node_dist = squared_distance(
107        point, median_node, axis_count=axis_count
108    )
109    # Update the best point and distance if necessary
110    if best is None or median_node_dist < best_dist:
111        best = median_node
112        best_dist = median_node_dist
113    # Calculate the difference for node selection given the current axis
114    axis = node[1]
115    diff = point[axis] - median_node[axis]
116    # Choose side to search
117    close, away = (node[2], node[3]) if diff < 0 else (node[3], node[2])
118    # Search the close side first
119    best, best_dist = closest_point(
120        close,
121        point,
122        best,
123        axis_count=axis_count,
124        best_dist=best_dist,
125    )
126    # Check the other side if needed
127    if diff**2 < best_dist:
128        best, best_dist = closest_point(
129            away,
130            point,
131            best,
132            axis_count=axis_count,
133            best_dist=best_dist,
134        )
135    return best, best_dist