Learning to Design Data-structures: A Case Study of Nearest Neighbor Search

We propose a general framework for automating data-structure design and apply it to the problem of nearest neighbor search. Our model adapts to the underlying data distribution and provides fine-grained control over query and space complexity, enabling the discovery of solutions tailored to problem-specific constraints. We are able to reverse-engineer learned algorithms in several settings. In 1D, the model discovers optimal distribution (in)dependent algorithms such as binary search and variants of interpolation search. In higher dimensions, the model learns solutions that resemble K-d trees in some regimes, while in others, have elements of locality-sensitive hashing.