k-d tree

Which town is nearest to the point where you are standing? The GeoNames gazetteer has 45,227 towns of the world with ten thousand people or more, and going through them all is 45 thousand distances per question. The tree below answers by looking at two or three dozen. Tap the map anywhere: the highlighted towns are the ones it looked at.

For Republic Square in Almaty the nearest town is, of course, Almaty, three kilometres away, and the tree needed 16 towns of 45,227. If the nearest neighbour is found this way for every town in the world, it comes to 28.6 towns per question on average. The loneliest is Nuuk: the nearest town, Keflavík, is 1404 km away. But from the North Pole the tree had to go through 1115 towns before making sure that the nearest is Hammerfest, 2150 km away. Why questions are cheap or expensive is below.

The idea

Take the town in the middle by longitude and draw a vertical line through it: the towns to the west go into the left subtree, to the east into the right. In each half take the town in the middle by latitude and cut it horizontally. Then by longitude again, and so on to the end. The result is a binary tree each node of which has a dividing line. To find the nearest town, go down to where the point lies, and on the way back look beyond a node’s line only if it is closer than the best town found.

This is the k-d tree, a k-dimensional search tree; in our case k = 2. It is a direct generalisation of the binary search tree: there the key is one and each node divides the line into “smaller” and “larger”, here there are two keys, and the nodes divide the plane by one and then the other in turn. The median at each node keeps the tree balanced: the halves are always equal, and the depth is log₂ n.

How it finds the nearest

Below are 237 towns around Almaty, from Samarkand to Kashgar, and their tree: each line divides its own rectangle, the one the node inherited from its ancestors. Press play to watch the search, or put the query point somewhere else.

The search starts at the root, as in an ordinary search tree: at each node compute the distance to its town, remember the best, and go down into the half where the query lies. On reaching a leaf we already have a candidate, usually a decent one. Then comes the main part. Going back up, at each node ask: can a town closer than the best lie beyond its line? Everything beyond the line is no closer to the query than the line itself. So if the distance from the query to the line is greater than the distance to the best town, that side can be thrown away whole, with all its towns. If less, it has to be looked into.

For Republic Square the search looked at 8 towns of 237, almost exactly the depth of the tree. The further the query is from all the towns, the larger the radius of the best and the more often its circle touches lines: for the point 42.0° N, 66.5° E the nearest town, Shardara, is 145 km away, and 14 towns had to be looked at. The North Pole in the first scene is an extreme case of the same thing: the best town is two thousand kilometres away, and almost every line is near.

The k nearest are found the same way: instead of one best, keep a heap of the k best with the worst of them on top, and compare the distance to the line with that worst. Switch k in the first scene.

The Earth is round: the third dimension

A world map has a catch: longitude and latitude are not flat coordinates. A degree of longitude is 111 km at the equator and zero at the pole, and a town “near in longitude” can be a thousand kilometres away. A tree on (longitude, latitude) would not find the nearest town. The way out is to go up a dimension. Write a point on the sphere as the unit vector (cos φ cos λ, cos φ sin λ, sin φ) and look for the nearest in a straight line through the Earth. The chord c grows together with the arc: the arc equals 2R·arcsin(c/2R), so the nearest by chord is also the nearest along the surface. And the chord is an ordinary Euclidean distance in three dimensions, and a k-d tree with k = 3 fits unchanged: the axes alternate x, y, z.

Code

A two-dimensional tree on the 237 towns around Almaty: building, the nearest town, the three nearest and a check by brute force. The coordinates are whole decimetres on a flat map, so all distances are exact.

And this is the first scene in numbers: a three-dimensional tree over the unit vectors of all 45,227 towns, the nearest neighbour for each of them and a few points at the edge of the world.

What it costs

Building: at each level find the median and split the points. With a sort at each level, as in the listings, this is O(n log² n); if the median is found by quickselect in linear time, O(n log n). Memory is a node per point, O(n).

A nearest-neighbour search for points scattered without particular structure visits O(log n) nodes on average: this was shown by Jerome Friedman, Jon Bentley and Raphael Finkel in 1977, while in the first paper on the tree Bentley estimated it by experiment. There is no worst-case guarantee: a distant query or points bunched on one line force a walk through most of the tree. It helps to keep at each node the rectangle containing its whole subtree and to compare the best distance with the distance to the rectangle rather than to the line; Friedman, Bentley and Finkel already proposed this check, and the problem below does it this way.

Brute force grows linearly with n, the tree barely noticeably: on 64 thousand points it is hundreds of times faster in Python and almost a hundred times faster in Go. Just do not forget the building: it costs more than one brute-force search, and if there is one query and the points are new each time, the tree will not pay off.

The curse of dimensionality

Pruning works as long as the circle around the query touches few lines. In many-dimensional space this stops being so: almost all points are at roughly the same distance from the query, the “ball of the best” is huge compared with the cells, and it touches almost all the lines. The listing builds a tree on 4096 random points in cubes of different dimensions and counts how many points the search looks at:

In two dimensions it is 17 points of 4096, in eight already 747, and in twelve 88% of all the points: the tree is no longer faster than brute force, and walking it costs time on top. Hence a well-known rule: a k-d tree is useful while there are many more points than 2^k, where k is the dimension. 4096 is exactly 2¹². For vectors of hundreds of dimensions, embeddings of texts and images, one looks not for the exact nearest but for a “close enough” one, with other structures.

Where it works

• Scientific libraries. scipy.spatial.KDTree in SciPy, neighbour search in scikit-learn: for points in two, three, a dozen dimensions.
• Ray tracing. A k-d tree over the scene’s triangles lets you check a ray’s intersection not with all the triangles but with the few that lie along its way.
• Maps. The nearest town, station, address to a point on the map; neighbours for “near you” suggestions. For “what falls into the map window” queries R-trees or space-filling curves are more often taken.
• Neighbour-based classification. The k nearest neighbours method and density-based clustering ask “who is near” for each data point; with a tree this is fast while there are few features.

History

Jon Bentley described the k-d tree in 1975 in Communications of the ACM, in the paper “Multidimensional binary search trees used for associative searching”. The paper showed that one and the same structure answers different queries, for example by part of the keys and the nearest-neighbour search, and noted from experiments that the time for the nearest-neighbour search grows as the logarithm of the number of points. Two years later Bentley, together with Jerome Friedman and Raphael Finkel, described nearest-neighbour search with logarithmic expected time. In that paper the axis of division at each node is the one along which the points are spread most, and one looks beyond the line only if the ball around the query touches the rectangle of the subtree.

Problems

In the problem you need the k nearest towns for thousands of queries, with ties and with towns lined up in a row, exactly where simple pruning by the line fails. One of the tests is the same 45 thousand towns of the world.