Learning Partition Trees for Nearest Neighbor Search

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.