Learning Partition Trees for Nearest Neighbor Search
2026-07-10 • Data Structures and Algorithms
Data Structures and AlgorithmsMachine Learning
AI summaryⓘ
The authors study how to make nearest neighbor search faster by learning a data structure based on the specific queries a user might ask. They focus on balanced halfspace trees, a method for dividing space to speed up searches, and show how to build such trees efficiently under certain assumptions about the data and query distribution. They prove that finding the best cut to make these trees is usually very hard, so they create a learning algorithm that finds an approximately good solution using more complex functions. This approach helps improve query times if a perfect tree structure exists for the data.
nearest neighbor searchdata-driven algorithm designbalanced halfspace treeslocality-sensitive hashingGaussian distributionspace partitioningNP-hardnesshalfspace cut problempolynomial threshold functionimproper learning
Authors
Sanjeev Khanna, Ashwin Padaki, Erik Waingarten
Abstract
We study nearest neighbor search from the perspective of data-driven algorithm design: given a dataset $P \subset \mathbb{R}^d$ of size $n$ and sample access to a query distribution over $\mathbb{R}^d$, the goal is to learn a data structure optimized for queries drawn from that specific distribution. We focus on the class of balanced halfspace trees, which naturally abstracts space-partitioning frameworks like locality-sensitive hashing. Assuming Gaussian-like marginal conditions on the dataset and query distribution, we give an efficient algorithm that learns a tree achieving $o(nd)$ query time, provided that a perfect tree exists. At the core of our algorithmic approach is the balanced halfspace cut problem, where we are given a distribution over $\mathbb{R}^d \times \mathbb{R}^d$ and must find a balanced halfspace that minimizes the fraction of cut pairs. We prove that without distributional assumptions, finding the optimal balanced halfspace is NP-hard. To circumvent this computational barrier, we design an efficient improper learning algorithm: if the optimal halfspace cuts an $α$ fraction of pairs, our algorithm outputs a balanced polynomial threshold function of degree $\tilde{O}(1/\varepsilon^2)$ that cuts at most an $O(\sqrt{α+\varepsilon})$ fraction.