cs.DSJul 10, 2026

Learning Partition Trees for Nearest Neighbor Search

Authors: Sanjeev KhannaAshwin PadakiErik Waingarten

Organizations: NYU · University of Pennsylvania

Abstract

We study nearest neighbor search from the perspective of data-driven algorithm design: given a dataset PRdP \subset \mathbb{R}^d of size nn and sample access to a query distribution over Rd\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)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 Rd×Rd\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 O~(1/ε2)\tilde{O}(1/\varepsilon^2) that cuts at most an O(α+ε)O(\sqrt{α+\varepsilon}) fraction.

Explore similar work

Sep 25, 2025cs.DS

Actively Learning Halfspaces without Synthetic Data

In the classic point location problem, one is given an arbitrary dataset XRdX \subset \mathbb{R}^d of nn points with query access to an unknown halfspace f:Rd{0,1}f : \mathbb{R}^d \to \{0,1\}, and the goal is to learn the label of every point in XX. This problem is extremely well-studied and a nearly-optimal O~(dlogn)\widetilde{O}(d \log n) query algorithm is known due to Hopkins-Kane-Lovett-Mahajan (FOCS 2020). However, their algorithm is granted the power to query arbitrary points outside of XX (point synthesis), and in fact without this power there is an Ω(n)Ω(n) query lower bound due to Dasgupta (NeurIPS 2004). In this work our goal is to design efficient algorithms for learning halfspaces without point synthesis. To circumvent the Ω(n)Ω(n) lower bound, we consider learning halfspaces whose normal vectors come from a set of size DD, and show tight bounds of Θ(D+logn)Θ(D + \log n). As a corollary, we obtain an optimal O(d+logn)O(d + \log n) query deterministic learner for axis-aligned halfspaces, closing a previous gap of O(dlogn)O(d \log n) vs. Ω(d+logn)Ω(d + \log n). In fact, our algorithm solves the more general problem of learning a Boolean function ff over nn elements which is monotone under at least one of DD provided orderings. Our technical insight is to exploit the structure in these orderings to perform a binary search in parallel rather than considering each ordering sequentially, and we believe our approach may be of broader interest. Furthermore, we use our exact learning algorithm to obtain nearly optimal algorithms for PAC-learning. We show that O(min(D+log(1/ε),1/ε)logD)O(\min(D + \log(1/\varepsilon), 1/\varepsilon) \cdot \log D) queries suffice to learn ff within error ε\varepsilon, even in a setting when ff can be adversarially corrupted on a cεc\varepsilon-fraction of points, for a sufficiently small constant cc. This bound is optimal up to a logD\log D factor, including in the realizable setting.
Hadley Black, Kasper Green Larsen, Arya Mazumdar +2
Aug 3, 2026cs.DS

Randomized Algorithms for Learning Partitions with Near Optimal Query Complexity in Constant Rounds

We study the round complexity of learning a hidden partition P\mathcal{P} of an nn-element universe using PAIR queries: PAIR(x,yx,y) tells us whether xx and yy belong to the same part of the partition or not. While it is easy to learn using nPn|\mathcal{P}| queries using a basic algorithm and this query complexity is optimal, this basic algorithm is highly sequential. Black, Mazumdar, and Saha [COLT 2025] recently gave tight deterministic round/query tradeoffs when the number of parts of P\mathcal{P} is known. In particular they prove Θ(loglogn)Θ(\log\log n) rounds are sufficient and necessary to limit the number of queries to nPn|\mathcal{P}|. They leave proving a randomized lower bound as an open direction. We show that randomization dramatically changes the picture. When the number of parts k=Pk = |\mathcal{P}| is known, we give a simple 3-round randomized algorithm using O(nklogn)O(nk\log n) queries with high probability, and prove that 2 rounds require Ω(n4/3k2/3)Ω(n^{4/3}k^{2/3}) queries -- the same as deterministic algorithms. We also study a more general setting where the number of parts is unknown. In this case, we give a 4-round randomized algorithm using O(nPlog2n)O(n|\mathcal P|\log^2 n) queries with high probability, and prove that 3-rounds cannot achieve near-optimal query complexity. Furthermore, we show an even bigger separation in this regime between randomized and deterministic algorithms: for the latter, Θ(logn/loglogn)Θ(\log n/\log\log n) rounds are necessary and sufficient to obtain near-optimal query complexity.
Deeparnab Chakrabarty, Aditi Dudeja, David Saulpic
Apr 16, 2026cs.DS

Tight Bounds for Learning Polyhedra with a Margin

We give an algorithm for PAC learning intersections of kk halfspaces with a ρρ margin to within error ε\varepsilon that runs in time poly(k,ε1,ρ1)exp(O(nlog(1/ρ)logk))\textsf{poly}(k, \varepsilon^{-1}, ρ^{-1}) \cdot \exp \left(O(\sqrt{n \log(1/ρ) \log k})\right). Notably, this improves on prior work which had an exponential dependence on either kk or ρ1ρ^{-1} and matches known cryptographic and Statistical Query lower bounds up to the logarithmic factors in kk and ρρ in the exponent. Our learning algorithm extends to the more general setting when we are only promised that most points have distance at least ρρ from the boundary of the polyhedron, making it applicable to continuous distributions as well.
Shyamal Patel, Santosh Vempala