Fast Diverse Nearest Neighbor Search

In the $k$-diverse nearest neighbor search problem, the input is a dataset in which each point is assigned a color, and given a query $q$, the goal is to retrieve $k$ approximate nearest neighbors with distinct colors. This color diversity metric is highly essential for information retrieval tasks where results are required to have different categories, e.g., recommending products from different sellers. Previous state-of-the-art solution by Anand et al. [ICML 2025] constructs a diversity-aware graph to meet this requirement. However, they suffer from an $\mathcal{O}(k^2)$ multiplicative query time overhead, creating a computational bottleneck when $k$ is large. To break this barrier, we propose a novel color bucketing framework that selects a subset of colors in each bucket and instantiates an independent approximate nearest neighbor search algorithm on the corresponding points. Inspired by group testing, we provide a randomized color sampling scheme that achieves an $\mathcal{O}\left(k^{1-\frac{1}{c^2}-o(1)}\log k\right)$ overhead in Euclidean spaces and an $\mathcal{O}\left(k \log k\right)$ overhead in general metric spaces, where $c$ is the approximation factor, reducing a factor of $k$. In addition, the space usage to store the data structure matches or improves upon prior methods. Empirical results on semi-synthetic and real-world datasets demonstrate that our approach achieves significantly faster search times and higher recall compared to state-of-the-art diversity search baselines. Furthermore, we introduce a dataset-oblivious deterministic bucketing scheme using expander graphs for a relaxed diverse search problem that only requires $k(1-\varepsilon)$ colors. Here, dataset-oblivious means that the color bucket construction does not depend on the spatial configuration of the dataset. We then establish an $\Omega(k^2)$ lower bound on the number of buckets for any dataset-oblivious deterministic bucketing scheme that gives an \emph{exact} $k$-diverse solution. We subsequently circumvent this by a dataset-aware adaptive segment tree bucketing scheme with an overhead of only $\mathcal{O}(k \log N)$, where $N$ is the number of colors. As an extension of our results, we solve the problem with general diversity metric, where the goal is to maximize the minimum pair-wise distance of the $k$ solutions, and improve the query time.