The KD-Tree

2 min read•1 page

A k-dimensional tree (KD-Tree) is a space-partitioning data structure for organizing points in a k-dimensional space. It is a binary search tree where every node represents a k-dimensional point.

KD-Tree Node:

1. Stores a point coordinate. 2. Stores the partition axis index. 3. Links to left (lesser) and right (greater) nodes.

python

1# k-Dimensional Node definition
2class KDNode:
3    def __init__(self, point, axis, left=None, right=None):
4        self.point = point     # k-dimensional coordinate (e.g. [x, y, z])
5        self.axis = axis       # Splitting axis index (0 for X, 1 for Y, etc.)
6        self.left = left       # Left subtree reference
7        self.right = right     # Right subtree reference

1 min read•1 page

To maintain a balanced tree, we partition the dataset at the median coordinate along the splitting axis.

Median Calculation:

1. Select the current axis (alternating X, Y, Z based on depth). 2. Sort points along this axis. 3. Choose the median point as the splitting pivot.

python

1# Select axis and find median point
2def get_median_pivot(points, axis):
3    # Sort points along the active axis
4    points_sorted = sorted(points, key=lambda p: p[axis])
5    median_idx = len(points_sorted) // 2
6    return points_sorted[median_idx]

Active Axis (0=X, 1=Y, 2=Z)

0.00

1 min read•1 page

A KD-tree is built recursively by alternating sort axis and splitting points at the median into left and right subtrees.

Tree Construction:

1. Select axis based on current depth. 2. Sort points and divide at median pivot. 3. Recurse for left and right point subsets.

python

1# Recursive KD-Tree Construction
2def build_kd_tree(points, depth=0):
3    if not points:
4        return None
5        
6    k = len(points[0])
7    axis = depth % k
8    
9    points.sort(key=lambda x: x[axis])
10    median = len(points) // 2
11    
12    return KDNode(
13        point=points[median],
14        axis=axis,
15        left=build_kd_tree(points[:median], depth + 1),
16        right=build_kd_tree(points[median + 1:], depth + 1)
17    )

Subdivision Depth

2.00

2 min read•1 page

To perform queries, the tree is traversed from the root, selecting sub-branches based on coordinate comparisons until a leaf cell is reached.

Leaf Traversal:

1. Start at the root node. 2. Compare target coordinate to node value along active axis. 3. Recurse down the appropriate branch until a leaf cell is reached.

python

1# Traverse down to leaf cell containing target point
2def find_leaf_cell(node, target, depth=0):
3    if node is None:
4        return None
5        
6    axis = depth % len(target)
7    if target[axis] < node.point[axis]:
8        if node.left is None:
9            return node
10        return find_leaf_cell(node.left, target, depth + 1)
11    else:
12        if node.right is None:
13            return node
14        return find_leaf_cell(node.right, target, depth + 1)

Target X

2.00

Target Y

2.00

1 min read•1 page

Once a leaf candidate is found, the search backtracks up the tree. We check if the distance to the splitting plane is smaller than the current best distance. If not, we prune the opposite branch to avoid redundant checks.

Backtracking Pruning:

1. Calculate distance from target to the node's splitting plane. 2. If the distance is smaller than current best, recursively search the opposite branch. 3. Otherwise, skip/prune the entire opposite subtree.

python

1# Backtrack and prune opposite branch if boundary is closer
2dist_to_boundary = (target[axis] - node.point[axis]) ** 2
3if dist_to_boundary < best_dist:
4    # Recurse down the opposite branch
5    best_point, best_dist = find_nearest(opposite_branch, target, depth + 1, best_point, best_dist)

Target X

4.00

Target Y

4.00

3 min read•1 page

To find the K closest points rather than just one, we use a max-priority queue (or heap) of size K to record the closest candidates found so far.

KNN Query:

1. Walk down the tree towards the target. 2. Insert visited points into a size-limited priority queue. 3. Backtrack, updating the pruning radius to the distance of the K-th closest candidate.

python

1# KNN query using a max-priority queue of size K
2def knn_search(node, target, k, heap=[], depth=0):
3    if node is None:
4        return
5        
6    dist = sum((a - b) ** 2 for a, b in zip(node.point, target))
7    
8    # Maintain heap of size K containing (negative_dist, point)
9    if len(heap) < k:
10        heapq.heappush(heap, (-dist, node.point))
11    elif dist < -heap[0][0]:
12        heapq.heapreplace(heap, (-dist, node.point))
13        
14    axis = depth % len(target)
15    next_node = node.left if target[axis] < node.point[axis] else node.right
16    other_node = node.right if target[axis] < node.point[axis] else node.left
17    
18    # Search subtree containing target first
19    knn_search(next_node, target, k, heap, depth + 1)
20    
21    # Check other subtree if coordinate distance is closer than current K-th best distance
22    dist_to_plane = (target[axis] - node.point[axis]) ** 2
23    if len(heap) < k or dist_to_plane < -heap[0][0]:
24        knn_search(other_node, target, k, heap, depth + 1)

Value of K

5.00

2 min read•1 page

To retrieve all points within a specific spatial box, we prune sub-branches that do not intersect the search bounding box.

Range Search:

1. Provide a bounding box defined by min and max bounds. 2. Traverse tree. If node lies in bounds, record it. 3. Compare splitting plane position to range bounds; traverse intersecting sub-branches only.

python

1# Range query finding all points within a bounding box
2def range_search(node, box_min, box_max, depth=0, results=[]):
3    if node is None:
4        return results
5        
6    # Check if current node is inside bounding box
7    if all(mn <= val <= mx for val, mn, mx in zip(node.point, box_min, box_max)):
8        results.append(node.point)
9        
10    axis = depth % len(box_min)
11    # Search left subtree if splitting plane is higher than box_min
12    if box_min[axis] <= node.point[axis]:
13        range_search(node.left, box_min, box_max, depth + 1, results)
14    # Search right subtree if splitting plane is lower than box_max
15    if box_max[axis] >= node.point[axis]:
16        range_search(node.right, box_min, box_max, depth + 1, results)
17        
18    return results

Box Size

6.00

2 min read•1 page

To find spatial overlap between two point clouds, we recursively test the bounding boxes of their KD-trees, avoiding costly all-to-all comparisons.

Tree-to-Tree Overlap:

1. Verify overlap between the current bounding volumes of Node A and Node B. 2. If they overlap, recursively check children sub-nodes. 3. If both are leaf nodes, perform direct point distance tests.

python

1# KD-Tree to KD-Tree Spatial Overlap Intersection
2def intersect_kd_trees(nodeA, nodeB, overlap_list):
3    if nodeA is None or nodeB is None:
4        return
5        
6    if not boxes_overlap(nodeA.bounds, nodeB.bounds):
7        return
8        
9    if nodeA.is_leaf and nodeB.is_leaf:
10        for pA in nodeA.points:
11            for pB in nodeB.points:
12                if distance(pA, pB) < tolerance:
13                    overlap_list.append((pA, pB))
14        return
15        
16    intersect_kd_trees(nodeA.left, nodeB.left, overlap_list)
17    intersect_kd_trees(nodeA.left, nodeB.right, overlap_list)
18    intersect_kd_trees(nodeA.right, nodeB.left, overlap_list)
19    intersect_kd_trees(nodeA.right, nodeB.right, overlap_list)

Bunny Offset Separation

6.00