Efficient K-Nearest Neighbors

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61

import numpy as np
from dataclasses import dataclass
from typing import Optional, List, Tuple
 
@dataclass
class KDNode:
    """
    A node in a KD-tree.
    
    Attributes:
    -----------
    point : np.ndarray
        The d-dimensional point stored at this node
    split_dim : int
        The dimension used for partitioning at this node
    split_value : float
        The splitting threshold (value of point[split_dim])
    left : Optional[KDNode]
        Left subtree (points with x[split_dim] <= split_value)
    right : Optional[KDNode]
        Right subtree (points with x[split_dim] > split_value)
    index : int
        Original index of this point in the dataset
    """
    point: np.ndarray
    split_dim: int
    split_value: float
    left: Optional['KDNode'] = None
    right: Optional['KDNode'] = None
    index: int = -1
    
    def is_leaf(self) -> bool:
        """Check if this node is a leaf (no children)."""
        return self.left is None and self.right is None
    
    @property
    def dimension(self) -> int:
        """Return the dimensionality of points in this tree."""
        return len(self.point)
 
 
@dataclass
class KDTree:
    """
    KD-Tree for efficient nearest neighbor search.
    
    Attributes:
    -----------
    root : Optional[KDNode]
        Root node of the tree
    n_points : int
        Total number of points in the tree
    n_dims : int
        Dimensionality of the data
    """
    root: Optional[KDNode]
    n_points: int
    n_dims: int
    
    def __repr__(self) -> str:
        return f"KDTree(n_points={self.n_points}, n_dims={self.n_dims})"

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111

import numpy as np
from typing import Optional, List
 
def build_kd_tree(
    points: np.ndarray,
    indices: Optional[np.ndarray] = None,
    depth: int = 0,
    split_strategy: str = "round_robin"
) -> Optional[KDNode]:
    """
    Build a KD-tree from a set of points.
    
    Time Complexity: O(n log n) with median selection
    Space Complexity: O(n) for the tree structure
    
    Parameters:
    -----------
    points : np.ndarray
        Array of points, shape (n, d)
    indices : np.ndarray, optional
        Original indices of points (for tracking)
    depth : int
        Current depth in the tree (for dimension cycling)
    split_strategy : str
        How to select split dimension: "round_robin", "max_variance", "max_extent"
    
    Returns:
    --------
    KDNode or None
        Root of the constructed (sub)tree
    """
    n, d = points.shape
    
    # Base case: no points
    if n == 0:
        return None
    
    # Initialize indices if not provided
    if indices is None:
        indices = np.arange(n)
    
    # Select splitting dimension
    if split_strategy == "round_robin":
        split_dim = depth % d
    elif split_strategy == "max_variance":
        # Split on dimension with maximum variance
        variances = np.var(points, axis=0)
        split_dim = int(np.argmax(variances))
    elif split_strategy == "max_extent":
        # Split on dimension with maximum range
        extents = np.ptp(points, axis=0)  # peak-to-peak
        split_dim = int(np.argmax(extents))
    else:
        split_dim = depth % d
    
    # Base case: single point
    if n == 1:
        return KDNode(
            point=points[0],
            split_dim=split_dim,
            split_value=points[0, split_dim],
            index=indices[0]
        )
    
    # Find median along split dimension
    # O(n) using selection algorithm, or O(n log n) using sort
    sorted_idx = np.argsort(points[:, split_dim])
    median_idx = n // 2
    
    # Handle even-sized arrays: median_idx is right-of-center
    # This ensures left subtree has <= ceil(n/2) - 1 points
    
    median_global_idx = sorted_idx[median_idx]
    median_point = points[median_global_idx]
    
    # Partition points
    left_mask = sorted_idx[:median_idx]
    right_mask = sorted_idx[median_idx + 1:]
    
    # Recursively build subtrees
    left_child = build_kd_tree(
        points[left_mask],
        indices[left_mask],
        depth + 1,
        split_strategy
    ) if len(left_mask) > 0 else None
    
    right_child = build_kd_tree(
        points[right_mask],
        indices[right_mask],
        depth + 1,
        split_strategy
    ) if len(right_mask) > 0 else None
    
    return KDNode(
        point=median_point,
        split_dim=split_dim,
        split_value=median_point[split_dim],
        left=left_child,
        right=right_child,
        index=indices[median_global_idx]
    )
 
 
def build_kd_tree_wrapper(points: np.ndarray) -> KDTree:
    """
    Public interface for KD-tree construction.
    """
    n, d = points.shape
    root = build_kd_tree(points)
    return KDTree(root=root, n_points=n, n_dims=d)

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145

import numpy as np
from typing import Tuple, Optional
import heapq
 
def nn_search(
    node: Optional[KDNode],
    query: np.ndarray,
    best: Tuple[float, Optional[KDNode]] = (float('inf'), None)
) -> Tuple[float, Optional[KDNode]]:
    """
    Find nearest neighbor in KD-tree using branch-and-bound.
    
    Time Complexity: O(log n) average, O(n) worst case
    
    Parameters:
    -----------
    node : KDNode
        Current node in the search
    query : np.ndarray
        Query point
    best : Tuple[float, KDNode]
        Current best (distance, node) pair
    
    Returns:
    --------
    Tuple[float, KDNode]
        Distance and node of nearest neighbor
    """
    if node is None:
        return best
    
    best_dist, best_node = best
    
    # Compute distance from query to current node's point
    curr_dist = np.sqrt(np.sum((node.point - query) ** 2))
    
    # Update best if current node is closer
    if curr_dist < best_dist:
        best_dist = curr_dist
        best_node = node
    
    # Determine which subtree to search first
    split_dim = node.split_dim
    split_val = node.split_value
    
    if query[split_dim] <= split_val:
        first, second = node.left, node.right
    else:
        first, second = node.right, node.left
    
    # Search the subtree containing the query first
    best_dist, best_node = nn_search(first, query, (best_dist, best_node))
    
    # Check if we need to search the other subtree
    # Distance from query to the splitting hyperplane
    dist_to_plane = abs(query[split_dim] - split_val)
    
    if dist_to_plane < best_dist:
        # The other subtree might contain closer points
        best_dist, best_node = nn_search(second, query, (best_dist, best_node))
    # else: PRUNED! The entire other subtree is guaranteed to be farther
    
    return best_dist, best_node
 
 
def knn_search(
    node: Optional[KDNode],
    query: np.ndarray,
    k: int,
    heap: list = None,
    depth: int = 0
) -> list:
    """
    Find K nearest neighbors in KD-tree.
    
    Uses a max-heap of size k to track the k closest points found so far.
    The max-heap allows O(1) access to the farthest of the k candidates,
    enabling efficient pruning.
    
    Parameters:
    -----------
    node : KDNode
        Current node in the search
    query : np.ndarray
        Query point
    k : int
        Number of neighbors to find
    heap : list
        Max-heap of (-distance, node) pairs
    
    Returns:
    --------
    list of (distance, node) pairs
        The k nearest neighbors
    """
    if heap is None:
        heap = []
    
    if node is None:
        return heap
    
    # Compute distance from query to current node's point
    dist = np.sqrt(np.sum((node.point - query) ** 2))
    
    # Update heap
    if len(heap) < k:
        # Heap not full, add directly
        heapq.heappush(heap, (-dist, id(node), node))
    elif dist < -heap[0][0]:
        # Current node closer than farthest in heap
        heapq.heapreplace(heap, (-dist, id(node), node))
    
    # Determine search order
    split_dim = node.split_dim
    split_val = node.split_value
    
    if query[split_dim] <= split_val:
        first, second = node.left, node.right
    else:
        first, second = node.right, node.left
    
    # Search first subtree
    knn_search(first, query, k, heap)
    
    # Pruning check: distance to splitting plane
    dist_to_plane = abs(query[split_dim] - split_val)
    
    # Only prune if heap is full AND plane is farther than k-th closest
    if len(heap) < k or dist_to_plane < -heap[0][0]:
        knn_search(second, query, k, heap)
    
    return heap
 
 
def query_kd_tree(tree: KDTree, query: np.ndarray, k: int = 1):
    """
    Public interface for KNN query.
    """
    if k == 1:
        dist, node = nn_search(tree.root, query)
        return [(dist, node.point, node.index)]
    else:
        heap = knn_search(tree.root, query, k)
        results = [(-d, node.point, node.index) for d, _, node in heap]
        return sorted(results)  # Sort by distance ascending

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55

@dataclass
class BucketKDNode:
    """
    KD-tree node that can store multiple points in leaves.
    """
    points: np.ndarray = None        # For leaf nodes: the actual points
    indices: np.ndarray = None       # Original indices of points
    split_dim: int = -1              # For internal nodes
    split_value: float = 0.0
    left: 'BucketKDNode' = None
    right: 'BucketKDNode' = None
    
    def is_leaf(self) -> bool:
        return self.points is not None
 
 
def build_bucket_kd_tree(
    points: np.ndarray,
    indices: np.ndarray = None,
    depth: int = 0,
    leaf_size: int = 20
) -> BucketKDNode:
    """
    Build KD-tree with bucket leaves for improved cache performance.
    """
    n, d = points.shape
    if indices is None:
        indices = np.arange(n)
    
    # Base case: create leaf if small enough
    if n <= leaf_size:
        return BucketKDNode(points=points.copy(), indices=indices.copy())
    
    # Split on cycling dimension
    split_dim = depth % d
    
    # Find median
    sorted_idx = np.argsort(points[:, split_dim])
    median_idx = n // 2
    split_value = points[sorted_idx[median_idx], split_dim]
    
    # Partition
    left_mask = sorted_idx[:median_idx]
    right_mask = sorted_idx[median_idx:]
    
    return BucketKDNode(
        split_dim=split_dim,
        split_value=split_value,
        left=build_bucket_kd_tree(
            points[left_mask], indices[left_mask], depth + 1, leaf_size
        ),
        right=build_bucket_kd_tree(
            points[right_mask], indices[right_mask], depth + 1, leaf_size
        )
    )