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 ℝd of size n and sample access to a query distribution over ℝ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 ℝd × ℝ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 tildeO(1/varepsilon2) that cuts at most an O(√α+varepsilon) fraction.