# Learning Partition Trees for Nearest Neighbor Search

arXiv: https://arxiv.org/abs/2607.09909 · HTML: https://mlsanity.com/paper/2607.09909/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 20 of its 48 references that are in the map.
- First posted on arXiv 10 Jul 2026.

## Builds on (most cited first, 20 of 20)

- [Efficient and robust approximate nearest neighbor search using Hierarchical Navigable Small World graphs](https://mlsanity.com/paper/1603.09320/) (2016)
- [Algorithms with Predictions](https://mlsanity.com/paper/2006.09123/) (2020)
- [A Survey on Learning to Hash](https://mlsanity.com/paper/1606.00185/) (2016)
- [A cost function for similarity-based hierarchical clustering](https://mlsanity.com/paper/1510.05043/) (2015)
- [A PAC Approach to Application-Specific Algorithm Selection](https://mlsanity.com/paper/1511.07147/) (2015)
- [Near-Optimal SQ Lower Bounds for Agnostically Learning Halfspaces and ReLUs under Gaussian Marginals](https://mlsanity.com/paper/2006.16200/) (2020)
- [A Faster Interior Point Method for Semidefinite Programming](https://mlsanity.com/paper/2009.10217/) (2020)
- [Approximate Nearest Neighbor Search in High Dimensions](https://mlsanity.com/paper/1806.09823/) (2018)
- [Results of the NeurIPS'21 Challenge on Billion-Scale Approximate Nearest Neighbor Search](https://mlsanity.com/paper/2205.03763/) (2022)
- [Optimal Data-Dependent Hashing for Approximate Near Neighbors](https://mlsanity.com/paper/1501.01062/) (2015)
- [Approximate Hierarchical Clustering via Sparsest Cut and Spreading Metrics](https://mlsanity.com/paper/1609.09548/) (2016)
- [Beyond Locality-Sensitive Hashing](https://mlsanity.com/paper/1306.1547/) (2013)
- [Worst-case Performance of Popular Approximate Nearest Neighbor Search Implementations: Guarantees and Limitations](https://mlsanity.com/paper/2310.19126/) (2023)
- [Hardness of Approximate Nearest Neighbor Search](https://mlsanity.com/paper/1803.00904/) (2018)
- [Optimal Hashing-based Time-Space Trade-offs for Approximate Near Neighbors](https://mlsanity.com/paper/1608.03580/) (2016)
- [Navigable Graphs for High-Dimensional Nearest Neighbor Search: Constructions and Limits](https://mlsanity.com/paper/2405.18680/) (2024)
- [Spectral Approaches to Nearest Neighbor Search](https://mlsanity.com/paper/1408.0751/) (2014)
- [Sparse Navigable Graphs for Nearest Neighbor Search: Algorithms and Hardness](https://mlsanity.com/paper/2507.14060/) (2025)
- [Data-Dependent LSH for the Earth Mover's Distance](https://mlsanity.com/paper/2403.05041/) (2024)
- [A Framework for Building Data Structures from Communication Protocols](https://mlsanity.com/paper/2506.20761/) (2025)

## Abstract (arXiv)

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.