Hi Friends,
Welcome to the 199th issue of the Polymathic Engineer newsletter. This week, we'll start talking about nearest neighbor search.
Suppose that you are building the checkout of a parcel-delivery network. The network includes pickup points (lockers) distributed across the country, and the application suggests the one closest to the specified delivery address.
With a few thousand lockers, this seems a trivial, quick problem. For each order, you compute the distance between the address and every locker and pick the smallest. A few thousand distances take a fraction of a millisecond, and a brute-force loop is the right solution.
Suppose now the network grows to hundreds of thousands of lockers, and every checkout waits on the answer. What you computed for the previous order doesn't help, because the address is different, so you start from scratch every time. This same type of problem arises in many other contexts, such as searching for “nearby restaurants” or “similar products.” The goal is always to identify, from a set of stored points, those closest to a given point.
In this issue, we will explore how to build a data structure capable of answering such queries without having to examine every single point. It is known as a k-d tree, a binary search tree with a twist.
The outline is as follows:
- Why the easy fixes don’t work
- One dimension is easy
- Why the obvious extension blows up
- Cycle through the dimensions
- Search and insert
- Nearest neighbor: descend, then backtrack
To learn technical skills, you must work on real projects. CodeCrafters is a great platform for that. You can build your own Redis, Kafka, DNS server, SQLite, HTTP server, or Git from scratch using your chosen programming language.
Why the easy fixes don’t work
The first idea is to precompute the closest locker for each customer, but this doesn’t really work. Customers can and will move, or they may want to ship to their home one day and to their office the next. You would need to maintain a list of lockers sorted by distance for each customer. This is one of those cases where caching doesn’t really help us a lot.
Thinking in 2-D, we might try a different approach inspired by how maps work: divide the map into tiles with a regular grid. To find which tile contains a point, we divide the value of each coordinate by the tile size, then compare the address only to lockers in the same tile and neighboring tiles. This cuts down the number of points we need to compare, but there is a catch. The approach only works if the points are distributed uniformly across the map, something that, with real-world data, typically does not happen. Lockers tend to be concentrated in urban centers and sparse in rural areas; as a result, some parts of the map may contain thousands of them, while most areas have none at all. We need a way of dividing the map that is more flexible, one that adapts to the actual distribution of the points.
When a problem seems too hard, it often helps to solve a simpler version first and then come up with a more general solution that works for the initial problem. So let’s restrict to one dimension.
One dimension is easy
Suppose that we only need to serve customers along a single road, and all our lockers are also placed along this same road. To simplify further, let’s assume the road is a straight line and that the total distance to be covered is short enough that the curvature of the Earth’s surface does not need to be taken into account. Every locker is then a distance from the start of the road, and a point on the map becomes a single number.
Given a random delivery address on the road, we want to know which locker is closest. If we keep the distances in a sorted array, this looks very similar to binary search. We start from the middle, compare the address with the element located there, and move either to the left or to the right until only two elements remain. If an exact match is found, we have identified the desired locker; otherwise, the closest locker will be one of the two elements adjacent to the position where the address would be inserted into the array.
However, arrays are not really flexible. If we would like to add a new locker between two existing ones, we have to shift every entry after it, and maybe reallocate the array if it’s static. Luckily, we know a data structure that is more flexible than arrays and still lets us perform binary search efficiently. A binary search tree (BST) is what we are looking for. Each node holds a key; the keys in its left subtree are smaller, and the keys in its right subtree are larger. The tree must be balanced to ensure logarithmic execution time for the most common operations. For a million lockers, that means about 20 comparisons instead of calculating a million distances. More importantly, the tree is a structure that can be generalized to two dimensions.
Finding the nearest locker in the tree works the same way as in the array. We search for the address, and if we find an exact match, that is the answer. Otherwise, the search ends at a node with a missing child, and the answer is either the node where the binary search fails or the last node where the search turned in the opposite direction. These are the two keys that bracket the address, the same two entries we compared in the array.
Now that we have solved the problem in one dimension, the question arises: can a similar data structure solve the problem in two dimensions?